fix: fix goroutine leaks around shutdown (#716)

Signed-off-by: Brandon Duffany <[email protected]>
Co-authored-by: Todd Baert <[email protected]>

Author: bduffany

providers/flagd/pkg/service/in_process/service.go

@@ -32,6 +32,7 @@ type InProcess struct {
 	serviceMetadata  model.Metadata
 	sync             sync.ISync
 	syncEnd          context.CancelFunc
+	wg               parallel.WaitGroup
 }
 
 type Configuration struct {
@@ -93,15 +94,19 @@ func (i *InProcess) Init() error {
 	syncChan := make(chan sync.DataSync, 1)
 
 	// start data sync
+	i.wg.Add(1)
 	go func() {
+		defer i.wg.Done()
 		err := i.sync.Sync(ctx, syncChan)
 		if err != nil {
 			syncInitErr <- err
 		}
 	}()
 
 	// start data sync listener and listen to listener shutdown hook
+	i.wg.Add(1)
 	go func() {
+		defer i.wg.Done()
 		for {
 			select {
 			case data := <-syncChan:
@@ -144,6 +149,7 @@ func (i *InProcess) Init() error {
 func (i *InProcess) Shutdown() {
 	i.syncEnd()
 	close(i.listenerShutdown)
+	i.wg.Wait()
 }
 
 func (i *InProcess) ResolveBoolean(ctx context.Context, key string, defaultValue bool,

providers/flagd/pkg/service/in_process/service_lifecycle_test.go

@@ -0,0 +1,74 @@
+package process
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+	"time"
+
+	of "github.com/open-feature/go-sdk/openfeature"
+)
+
+func TestInProcessServiceShutdownCleansUpGoroutines(t *testing.T) {
+	checkGoroutineLeaks(t)
+
+	flagFile := "config.json"
+	offlinePath := filepath.Join(t.TempDir(), flagFile)
+	if err := os.WriteFile(offlinePath, []byte(flagRsp), 0644); err != nil {
+		t.Fatal(err)
+	}
+	service := NewInProcessService(Configuration{OfflineFlagSource: offlinePath})
+	if err := service.Init(); err != nil {
+		t.Fatal(err)
+	}
+
+	// Wait for provider to be ready.
+	channel := service.EventChannel()
+	select {
+	case event := <-channel:
+		if event.EventType != of.ProviderReady {
+			t.Fatalf("Provider initialization failed. Got event type %s with message %s", event.EventType, event.Message)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Provider initialization did not complete within acceptable timeframe")
+	}
+
+	// Service is ready - now shut down.
+	service.Shutdown()
+}
+
+func TestInProcessServiceImmediateShutdownCleansUpGoroutines(t *testing.T) {
+	checkGoroutineLeaks(t)
+
+	flagFile := "config.json"
+	offlinePath := filepath.Join(t.TempDir(), flagFile)
+	if err := os.WriteFile(offlinePath, []byte(flagRsp), 0644); err != nil {
+		t.Fatal(err)
+	}
+	service := NewInProcessService(Configuration{OfflineFlagSource: offlinePath})
+	if err := service.Init(); err != nil {
+		t.Fatal(err)
+	}
+
+	// Immediately shut down - don't wait for the provider to be ready.
+	service.Shutdown()
+}
+
+// At the end of the test, if no other failures have occurred, check for
+// goroutine leaks, and fail the test if any were found.
+func checkGoroutineLeaks(t *testing.T) {
+	startingGoroutineCount := runtime.NumGoroutine()
+	t.Cleanup(func() {
+		if t.Failed() {
+			return
+		}
+		buf := make([]byte, 1<<20)
+		stacklen := runtime.Stack(buf, true)
+		if numGoroutinesAfter := runtime.NumGoroutine(); numGoroutinesAfter > startingGoroutineCount {
+			t.Errorf("Goroutines leaked: %d goroutines before, %d goroutines after", startingGoroutineCount, numGoroutinesAfter)
+			fmt.Fprintf(os.Stderr, "%s\n", buf[:stacklen])
+		}
+	})
+}

providers/flagd/pkg/service/rpc/service.go

@@ -8,6 +8,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"sync"
 	"time"
 
 	schemaConnectV1 "buf.build/gen/go/open-feature/flagd/connectrpc/go/flagd/evaluation/v1/evaluationv1connect"
@@ -51,6 +52,7 @@ type Service struct {
 
 	client     schemaConnectV1.ServiceClient
 	cancelHook context.CancelFunc
+	wg         sync.WaitGroup
 }
 
 func NewService(cfg Configuration, cache *cache.Service, logger logr.Logger, retries int) *Service {
@@ -86,7 +88,9 @@ func (s *Service) Init() error {
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	s.cancelHook = cancelFunc
 
+	s.wg.Add(1)
 	go func() {
+		defer s.wg.Done()
 		s.startEventStream(ctx)
 	}()
 
@@ -97,6 +101,7 @@ func (s *Service) Shutdown() {
 	if s.cancelHook != nil {
 		s.cancelHook()
 	}
+	s.wg.Wait()
 }
 
 // ResolveBoolean handles the flag evaluation response from the flagd ResolveBoolean rpc
@@ -450,24 +455,28 @@ func (s *Service) startEventStream(ctx context.Context) {
 			}
 
 			// error in stream handler, purge cache if available and retry
-			s.logger.V(logger.Warn).Info("connection to event stream failed, attempting again")
+			s.logger.V(logger.Warn).Info(fmt.Sprintf("connection to event stream failed (%q), attempting again", err))
 			if s.cache.IsEnabled() {
 				s.cache.GetCache().Purge()
 			}
 		}
 
-		time.Sleep(s.retryCounter.sleep())
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(s.retryCounter.sleep()):
+		}
 	}
 
 	// retry attempts exhausted. Disable cache and emit error event
 	s.cache.Disable()
-	s.events <- of.Event{
+	s.sendEvent(ctx, of.Event{
 		ProviderName: "flagd",
 		EventType:    of.ProviderError,
 		ProviderEventDetails: of.ProviderEventDetails{
 			Message: "grpc connection establishment failed",
 		},
-	}
+	})
 }
 
 // streamClient opens the event stream and distribute streams to appropriate handlers.
@@ -485,9 +494,9 @@ func (s *Service) streamClient(ctx context.Context) error {
 
 		switch stream.Msg().Type {
 		case string(flagdService.ConfigurationChange):
-			s.handleConfigurationChangeEvent(stream.Msg())
+			s.handleConfigurationChangeEvent(ctx, stream.Msg())
 		case string(flagdService.ProviderReady):
-			s.handleReadyEvent()
+			s.handleReadyEvent(ctx)
 		case string(flagdService.Shutdown):
 			// this is considered as a non-error
 			return nil
@@ -504,7 +513,7 @@ func (s *Service) streamClient(ctx context.Context) error {
 	return nil
 }
 
-func (s *Service) handleConfigurationChangeEvent(event *schemaV1.EventStreamResponse) {
+func (s *Service) handleConfigurationChangeEvent(ctx context.Context, event *schemaV1.EventStreamResponse) {
 	if !s.cache.IsEnabled() {
 		return
 	}
@@ -536,20 +545,28 @@ func (s *Service) handleConfigurationChangeEvent(event *schemaV1.EventStreamResp
 		keys = append(keys, flagKey)
 	}
 
-	s.events <- of.Event{
+	s.sendEvent(ctx, of.Event{
 		ProviderName: "flagd",
 		EventType:    of.ProviderConfigChange,
 		ProviderEventDetails: of.ProviderEventDetails{
 			Message:     "flags changed",
 			FlagChanges: keys,
 		},
-	}
+	})
 }
 
-func (s *Service) handleReadyEvent() {
-	s.events <- of.Event{
+func (s *Service) handleReadyEvent(ctx context.Context) {
+	s.sendEvent(ctx, of.Event{
 		ProviderName: "flagd",
 		EventType:    of.ProviderReady,
+	})
+}
+
+func (s *Service) sendEvent(ctx context.Context, event of.Event) {
+	select {
+	case <-ctx.Done():
+		return
+	case s.events <- event:
 	}
 }
 

providers/flagd/pkg/service/rpc/service_eventing_test.go

@@ -1,14 +1,15 @@
 package rpc
 
 import (
-	schemaV1 "buf.build/gen/go/open-feature/flagd/protocolbuffers/go/flagd/evaluation/v1"
 	"context"
 	"errors"
+	"testing"
+	"time"
+
+	schemaV1 "buf.build/gen/go/open-feature/flagd/protocolbuffers/go/flagd/evaluation/v1"
 	"github.com/open-feature/go-sdk-contrib/providers/flagd/internal/cache"
 	of "github.com/open-feature/go-sdk/openfeature"
 	"google.golang.org/protobuf/types/known/structpb"
-	"testing"
-	"time"
 )
 
 func TestRetries(t *testing.T) {
@@ -69,7 +70,8 @@ func TestConfigChange(t *testing.T) {
 
 		// when
 		go func() {
-			service.handleConfigurationChangeEvent(&schemaV1.EventStreamResponse{
+			ctx := context.Background()
+			service.handleConfigurationChangeEvent(ctx, &schemaV1.EventStreamResponse{
 				Data: stData,
 			})
 		}()
@@ -93,7 +95,8 @@ func TestConfigChange(t *testing.T) {
 
 		// when
 		go func() {
-			service.handleConfigurationChangeEvent(&schemaV1.EventStreamResponse{
+			ctx := context.Background()
+			service.handleConfigurationChangeEvent(ctx, &schemaV1.EventStreamResponse{
 				Data: stData,
 			})
 		}()