Merge pull request #716 from strukturag/leak-grpc-goroutines

Prevent goroutine leaks in GRPC tests.

certificate_reloader.go

@@ -66,6 +66,11 @@ func NewCertificateReloader(certFile string, keyFile string) (*CertificateReload
 	return reloader, nil
 }
 
+func (r *CertificateReloader) Close() {
+	r.keyWatcher.Close()
+	r.certWatcher.Close()
+}
+
 func (r *CertificateReloader) reload(filename string) {
 	log.Printf("reloading certificate from %s with %s", r.certFile, r.keyFile)
 	pair, err := tls.LoadX509KeyPair(r.certFile, r.keyFile)
@@ -135,6 +140,10 @@ func NewCertPoolReloader(certFile string) (*CertPoolReloader, error) {
 	return reloader, nil
 }
 
+func (r *CertPoolReloader) Close() {
+	r.certWatcher.Close()
+}
+
 func (r *CertPoolReloader) reload(filename string) {
 	log.Printf("reloading certificate pool from %s", r.certFile)
 	pool, err := loadCertPool(r.certFile)

file_watcher.go

@@ -22,6 +22,7 @@
 package signaling
 
 import (
+	"context"
 	"errors"
 	"log"
 	"os"
@@ -55,6 +56,8 @@ type FileWatcher struct {
 	callback FileWatcherCallback
 
 	watcher   *fsnotify.Watcher
+	closeCtx  context.Context
+	closeFunc context.CancelFunc
 }
 
 func NewFileWatcher(filename string, callback FileWatcherCallback) (*FileWatcher, error) {
@@ -78,17 +81,23 @@ func NewFileWatcher(filename string, callback FileWatcherCallback) (*FileWatcher
 		return nil, err
 	}
 
+	closeCtx, closeFunc := context.WithCancel(context.Background())
+
 	w := &FileWatcher{
 		filename: filename,
 		target:   realFilename,
 		callback: callback,
 		watcher:  watcher,
+
+		closeCtx:  closeCtx,
+		closeFunc: closeFunc,
 	}
 	go w.run()
 	return w, nil
 }
 
 func (f *FileWatcher) Close() error {
+	f.closeFunc()
 	return f.watcher.Close()
 }
 
@@ -152,6 +161,8 @@ func (f *FileWatcher) run() {
 			}
 
 			log.Printf("Error watching %s: %s", f.filename, err)
+		case <-f.closeCtx.Done():
+			return
 		}
 	}
 }

file_watcher_test.go

@@ -47,6 +47,7 @@ func TestFileWatcher_NotExist(t *testing.T) {
 }
 
 func TestFileWatcher_File(t *testing.T) {
+	ensureNoGoroutinesLeak(t, func(t *testing.T) {
 		tmpdir := t.TempDir()
 		filename := path.Join(tmpdir, "test.txt")
 		if err := os.WriteFile(filename, []byte("Hello world!"), 0644); err != nil {
@@ -89,6 +90,7 @@ func TestFileWatcher_File(t *testing.T) {
 			t.Error("should not have received another event")
 		case <-ctxTimeout.Done():
 		}
+	})
 }
 
 func TestFileWatcher_Rename(t *testing.T) {

grpc_client.go

@@ -38,6 +38,7 @@ import (
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"google.golang.org/grpc"
 	codes "google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/resolver"
 	status "google.golang.org/grpc/status"
 )
@@ -275,6 +276,7 @@ type GrpcClients struct {
 	targetPrefix      string
 	targetInformation map[string]*GrpcTargetInformationEtcd
 	dialOptions       atomic.Value // []grpc.DialOption
+	creds             credentials.TransportCredentials
 
 	initializedCtx       context.Context
 	initializedFunc      context.CancelFunc
@@ -308,6 +310,13 @@ func (c *GrpcClients) load(config *goconf.ConfigFile, fromReload bool) error {
 		return err
 	}
 
+	if c.creds != nil {
+		if cr, ok := c.creds.(*reloadableCredentials); ok {
+			cr.Close()
+		}
+	}
+	c.creds = creds
+
 	opts := []grpc.DialOption{grpc.WithTransportCredentials(creds)}
 	c.dialOptions.Store(opts)
 
@@ -490,6 +499,7 @@ func (c *GrpcClients) loadTargetsStatic(config *goconf.ConfigFile, fromReload bo
 		entry, found := clientsMap[target]
 		if !found {
 			entry = &grpcClientsList{}
+			clientsMap[target] = entry
 		}
 		entry.clients = append(entry.clients, client)
 		clients = append(clients, client)
@@ -794,6 +804,11 @@ func (c *GrpcClients) Close() {
 	if c.etcdClient != nil {
 		c.etcdClient.RemoveListener(c)
 	}
+	if c.creds != nil {
+		if cr, ok := c.creds.(*reloadableCredentials); ok {
+			cr.Close()
+		}
+	}
 	c.closeFunc()
 }
 

grpc_client_test.go

@@ -112,6 +112,7 @@ func waitForEvent(ctx context.Context, t *testing.T, ch <-chan struct{}) {
 }
 
 func Test_GrpcClients_EtcdInitial(t *testing.T) {
+	ensureNoGoroutinesLeak(t, func(t *testing.T) {
 		_, addr1 := NewGrpcServerForTest(t)
 		_, addr2 := NewGrpcServerForTest(t)
 
@@ -130,6 +131,7 @@ func Test_GrpcClients_EtcdInitial(t *testing.T) {
 		if clients := client.GetClients(); len(clients) != 2 {
 			t.Errorf("Expected two clients, got %+v", clients)
 		}
+	})
 }
 
 func Test_GrpcClients_EtcdUpdate(t *testing.T) {
@@ -231,6 +233,7 @@ func Test_GrpcClients_EtcdIgnoreSelf(t *testing.T) {
 }
 
 func Test_GrpcClients_DnsDiscovery(t *testing.T) {
+	ensureNoGoroutinesLeak(t, func(t *testing.T) {
 		lookup := newMockDnsLookupForTest(t)
 		target := "testgrpc:12345"
 		ip1 := net.ParseIP("192.168.0.1")
@@ -282,6 +285,7 @@ func Test_GrpcClients_DnsDiscovery(t *testing.T) {
 		} else if !clients[0].ip.Equal(ip2) {
 			t.Errorf("Expected IP %s, got %s", ip2, clients[0].ip)
 		}
+	})
 }
 
 func Test_GrpcClients_DnsDiscoveryInitialFailed(t *testing.T) {
@@ -320,6 +324,7 @@ func Test_GrpcClients_DnsDiscoveryInitialFailed(t *testing.T) {
 }
 
 func Test_GrpcClients_Encryption(t *testing.T) {
+	ensureNoGoroutinesLeak(t, func(t *testing.T) {
 		serverKey, err := rsa.GenerateKey(rand.Reader, 1024)
 		if err != nil {
 			t.Fatal(err)
@@ -371,4 +376,5 @@ func Test_GrpcClients_Encryption(t *testing.T) {
 				t.Fatal(err)
 			}
 		}
+	})
 }

grpc_common.go

@@ -125,6 +125,15 @@ func (c *reloadableCredentials) OverrideServerName(serverName string) error {
 	return nil
 }
 
+func (c *reloadableCredentials) Close() {
+	if c.loader != nil {
+		c.loader.Close()
+	}
+	if c.pool != nil {
+		c.pool.Close()
+	}
+}
+
 func NewReloadableCredentials(config *goconf.ConfigFile, server bool) (credentials.TransportCredentials, error) {
 	var prefix string
 	var caPrefix string

grpc_server.go

@@ -108,6 +108,9 @@ func (s *GrpcServer) Run() error {
 
 func (s *GrpcServer) Close() {
 	s.conn.GracefulStop()
+	if cr, ok := s.creds.(*reloadableCredentials); ok {
+		cr.Close()
+	}
 }
 
 func (s *GrpcServer) LookupSessionId(ctx context.Context, request *LookupSessionIdRequest) (*LookupSessionIdReply, error) {