added manger to keep track of go routines in the services (#2869)

- added manager to wait for all go routines to end before exit
- code refactor
- renamed Manafer to Interface and GoRoutineManager to GroupManager
- replaced some go routine calls with manager Add func
- added unit tests for manager

pkg/utils/httpserver/server.go

@@ -6,10 +6,11 @@ import (
 	"net/http"
 	"strings"
 
+	"github.com/fission/fission/pkg/utils/manager"
 	"go.uber.org/zap"
 )
 
-func StartServer(ctx context.Context, log *zap.Logger, svc string, port string, handler http.Handler) {
+func StartServer(ctx context.Context, log *zap.Logger, mgr manager.Interface, svc string, port string, handler http.Handler) {
 	if !strings.Contains(port, ":") {
 		port = fmt.Sprintf(":%s", port)
 	}
@@ -19,13 +20,13 @@ func StartServer(ctx context.Context, log *zap.Logger, svc string, port string,
 	}
 	l := log.With(zap.String("service", svc), zap.String("addr", server.Addr))
 	l.Info("starting server")
-	go func() {
+	mgr.Add(ctx, func(ctx context.Context) {
 		if err := server.ListenAndServe(); err != nil {
 			if err != http.ErrServerClosed {
 				l.Error("server error", zap.Error(err))
 			}
 		}
-	}()
+	})
 	<-ctx.Done()
 	l.Info("shutting down server")
 	if err := server.Shutdown(ctx); err != nil {

pkg/utils/httpserver/server_test.go

@@ -12,9 +12,13 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/fission/fission/pkg/utils/loggerfactory"
+	"github.com/fission/fission/pkg/utils/manager"
 )
 
 func TestStartServer(t *testing.T) {
+	mgr := manager.New()
+	defer mgr.Wait()
+
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	logger := loggerfactory.GetLogger()
@@ -26,7 +30,10 @@ func TestStartServer(t *testing.T) {
 			logger.Error("failed to write response", zap.Error(err))
 		}
 	}))
-	go StartServer(ctx, logger, "test", "8999", m)
+
+	mgr.Add(ctx, func(ctx context.Context) {
+		StartServer(ctx, logger, mgr, "test", "8999", m)
+	})
 
 	tests := []struct {
 		Name       string

pkg/utils/manager/manager.go

@@ -0,0 +1,62 @@
+package manager
+
+import (
+	"context"
+	"sync"
+	"time"
+)
+
+var _ Interface = &GroupManager{}
+
+// Interface keeps track of the go routines in the system and can be used to gracefully shutdown the
+// the system by waiting for completion of go routines added to it.
+type Interface interface {
+	// Add will start a go routine for the given "function" and adds it to the list of go routines
+	// and will also remove the "function" from the list when it completes
+	Add(ctx context.Context, function func(context.Context))
+
+	// Wait blocks the execution of the process until all the go routines in the manager are completed.
+	Wait()
+
+	// WaitWithTimeout blocks the execution of the process until timeout or till all all the go routines in the manager are completed
+	WaitWithTimeout(timeout time.Duration) error
+}
+
+type GroupManager struct {
+	wg sync.WaitGroup
+}
+
+func New() Interface {
+	return &GroupManager{
+		wg: sync.WaitGroup{},
+	}
+}
+func (g *GroupManager) Add(ctx context.Context, f func(context.Context)) {
+	g.wg.Add(1)
+	go func() {
+		defer g.wg.Done()
+		f(ctx)
+	}()
+}
+
+func (g *GroupManager) Wait() {
+	g.wg.Wait()
+}
+
+func (g *GroupManager) WaitWithTimeout(timeout time.Duration) error {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	done := make(chan struct{})
+	go func() {
+		g.wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-done:
+		return nil
+	}
+}

pkg/utils/manager/manager_test.go

@@ -0,0 +1,65 @@
+package manager
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestAddAndWait(t *testing.T) {
+	mgr := New()
+
+	value := 10
+	expectedValue := 11
+
+	mgr.Add(context.Background(), func(ctx context.Context) {
+		time.Sleep(1 * time.Second)
+		value = expectedValue
+	})
+
+	mgr.Wait()
+	require.Equal(t, expectedValue, value, "manager did not wait for go routine to complete")
+}
+
+func TestAddWithContextCancel(t *testing.T) {
+	mgr := New()
+
+	value := 10
+	expectedValue := 11
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	mgr.Add(ctx, func(ctx context.Context) {
+		<-ctx.Done()
+		value = 11
+	})
+
+	go cancel()
+	mgr.Wait()
+	require.Equal(t, expectedValue, value, "manager did not wait for go routine to complete when context is cancelled")
+}
+
+func TestAddAndWaitWithTimeout(t *testing.T) {
+	mgr := New()
+
+	value := 10
+
+	mgr.Add(context.Background(), func(ctx context.Context) {
+		time.Sleep(2 * time.Second)
+	})
+
+	err := mgr.WaitWithTimeout(1 * time.Second)
+	require.NotNil(t, err, "manager WaitWithTimeout did not return an error when timeout exceeded")
+
+	expectedValue := 11
+	mgr.Add(context.Background(), func(ctx context.Context) {
+		time.Sleep(100 * time.Millisecond)
+		value = expectedValue
+	})
+
+	err = mgr.WaitWithTimeout(1 * time.Second)
+	require.Nil(t, err, "manager returned error even though all go routins completed successfully before timeout")
+	require.Equal(t, expectedValue, value)
+}

pkg/utils/metrics/server.go

@@ -26,9 +26,10 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/metrics"
 
 	"github.com/fission/fission/pkg/utils/httpserver"
+	"github.com/fission/fission/pkg/utils/manager"
 )
 
-func ServeMetrics(ctx context.Context, parent string, logger *zap.Logger) {
+func ServeMetrics(ctx context.Context, parent string, logger *zap.Logger, mgr manager.Interface) {
 	metricsAddr := os.Getenv("METRICS_ADDR")
 	if metricsAddr == "" {
 		metricsAddr = "8080"
@@ -45,5 +46,5 @@ func ServeMetrics(ctx context.Context, parent string, logger *zap.Logger) {
 			EnableOpenMetrics: true,
 		},
 	))
-	httpserver.StartServer(ctx, logger, parent+"/metrics", metricsAddr, mux)
+	httpserver.StartServer(ctx, logger, mgr, parent+"/metrics", metricsAddr, mux)
 }

pkg/utils/profile/profile.go

@@ -32,9 +32,10 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/fission/fission/pkg/utils/httpserver"
+	"github.com/fission/fission/pkg/utils/manager"
 )
 
-func ProfileIfEnabled(ctx context.Context, logger *zap.Logger) {
+func ProfileIfEnabled(ctx context.Context, logger *zap.Logger, mgr manager.Interface) {
 	enablePprof := os.Getenv("PPROF_ENABLED")
 	if enablePprof != "true" {
 		return
@@ -47,5 +48,7 @@ func ProfileIfEnabled(ctx context.Context, logger *zap.Logger) {
 	pprofMux := http.DefaultServeMux
 	http.DefaultServeMux = http.NewServeMux()
 
-	go httpserver.StartServer(ctx, logger, "pprof", pprofPort, pprofMux)
+	mgr.Add(ctx, func(ctx context.Context) {
+		httpserver.StartServer(ctx, logger, mgr, "pprof", pprofPort, pprofMux)
+	})
 }