Merge pull request #32 from did-method-plc/DavidBuchanan314/hang-fix

replica: Fix context handling and replica shutdown

Author: DavidBuchanan314

cmd/plc-replica/main.go

@@ -137,13 +137,24 @@ func run(ctx context.Context, cmd *cli.Command) error {
 	server := replica.NewServer(store, state, httpAddr, logger)
 	g, gctx := errgroup.WithContext(ctx)
 
-	g.Go(server.Run)
+	g.Go(func() error {
+		return server.Run(gctx)
+	})
 
 	g.Go(func() error {
 		mux := http.NewServeMux()
 		mux.Handle("/metrics", promhttp.Handler())
+		srv := &http.Server{Addr: metricsAddr, Handler: mux}
+		go func() {
+			<-gctx.Done()
+			srv.Shutdown(context.Background())
+		}()
 		slog.Info("metrics server listening", "addr", metricsAddr)
-		return http.ListenAndServe(metricsAddr, mux)
+		err := srv.ListenAndServe()
+		if err == http.ErrServerClosed {
+			return nil
+		}
+		return err
 	})
 
 	if !noIngest {

replica/ingest.go

@@ -196,7 +196,10 @@ func (i *Ingestor) Run(ctx context.Context) error {
 	}()
 
 	// Start ingestion state machine in a goroutine
-	go i.ingestLoop(ctx, &cursor, ingestedOps)
+	go func() {
+		i.ingestLoop(ctx, &cursor, ingestedOps)
+		close(ingestedOps)
+	}()
 
 	// Process operations from ingestion channel and add to InFlight before sending to workers
 	for seqop := range ingestedOps {

replica/server.go

@@ -1,6 +1,7 @@
 package replica
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"log/slog"
@@ -54,8 +55,8 @@ func NewServer(store *GormOpStore, state *ReplicaState, addr string, logger *slo
 	}
 }
 
-// Run starts the HTTP server (blocking)
-func (s *Server) Run() error {
+// Run starts the HTTP server (blocking). It shuts down gracefully when ctx is cancelled.
+func (s *Server) Run(ctx context.Context) error {
 	mux := http.NewServeMux()
 	mux.HandleFunc("GET /_health", s.handleHealth)
 	mux.HandleFunc("GET /{did}/log/audit", s.handleDIDLogAudit)
@@ -67,8 +68,18 @@ func (s *Server) Run() error {
 
 	handler := otelhttp.NewHandler(mux, "")
 
+	srv := &http.Server{Addr: s.addr, Handler: handler}
+	go func() {
+		<-ctx.Done()
+		srv.Shutdown(context.Background())
+	}()
+
 	s.logger.Info("http server listening", "addr", s.addr)
-	return http.ListenAndServe(s.addr, handler)
+	err := srv.ListenAndServe()
+	if err == http.ErrServerClosed {
+		return nil
+	}
+	return err
 }
 
 // handleIndex serves the index page