diff --git a/core/pkg/service/sync/handler.go b/core/pkg/service/sync/handler.go index 96defcb4c..f1b997b56 100644 --- a/core/pkg/service/sync/handler.go +++ b/core/pkg/service/sync/handler.go @@ -17,6 +17,8 @@ type handler struct { syncv1grpc.UnimplementedFlagSyncServiceServer syncStore subscriptions.Manager logger *logger.Logger + // ctx is used to handle SIG[INT|TERM] + ctx context.Context } func (nh *handler) SyncFlags( @@ -40,6 +42,8 @@ func (nh *handler) SyncFlags( } case <-server.Context().Done(): return nil + case <-nh.ctx.Done(): + return nil } } } @@ -64,6 +68,8 @@ type oldHandler struct { rpc.UnimplementedFlagSyncServiceServer syncStore subscriptions.Manager logger *logger.Logger + // ctx is used to handle SIG[INT|TERM] + ctx context.Context } func (l *oldHandler) FetchAllFlags(ctx context.Context, req *syncv1.FetchAllFlagsRequest) ( @@ -102,6 +108,8 @@ func (l *oldHandler) SyncFlags( } case <-stream.Context().Done(): return nil + case <-l.ctx.Done(): + return nil } } } diff --git a/core/pkg/service/sync/server.go b/core/pkg/service/sync/server.go index 068eead23..cf30d7f1d 100644 --- a/core/pkg/service/sync/server.go +++ b/core/pkg/service/sync/server.go @@ -35,14 +35,16 @@ type Server struct { metricServerReady bool } -func NewServer(logger *logger.Logger, store subscriptions.Manager) *Server { +func NewServer(ctx context.Context, logger *logger.Logger, store subscriptions.Manager) *Server { theOldHandler := &oldHandler{ logger: logger, syncStore: store, + ctx: ctx, } theNewHandler := &handler{ logger: logger, syncStore: store, + ctx: ctx, } return &Server{ oldHandler: theOldHandler, diff --git a/flagd-proxy/cmd/start.go b/flagd-proxy/cmd/start.go index 8c0be5145..9e8ee64b9 100644 --- a/flagd-proxy/cmd/start.go +++ b/flagd-proxy/cmd/start.go @@ -63,7 +63,7 @@ var startCmd = &cobra.Command{ ctx, _ := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) syncStore := subscriptions.NewManager(ctx, logger) - s := syncServer.NewServer(logger, syncStore) + s := syncServer.NewServer(ctx, logger, syncStore) cfg := service.Configuration{ ReadinessProbe: func() bool { return true },