Fixed operator signal interruption using context WithCancel

Signed-off-by: Tanton <[email protected]>

Author: Tanton

cmd/aws-service-operator/server.go

@@ -1,13 +1,15 @@
 package main
 
 import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+
 	"github.com/awslabs/aws-service-operator/pkg/logger"
 	"github.com/awslabs/aws-service-operator/pkg/server"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
-	"os"
-	"os/signal"
-	"syscall"
 )
 
 var serverCmd = &cobra.Command{
@@ -26,21 +28,15 @@ var serverCmd = &cobra.Command{
 		}
 		config.Logger = logger
 
+		ctx, cancel := context.WithCancel(context.Background())
 		signalChan := make(chan os.Signal, 1)
-		stopChan := make(chan struct{})
 		signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
 
-		server.New(config).Run(stopChan)
-
-		for {
-			select {
-			case <-signalChan:
-				logger.Info("shutdown signal received, exiting...")
-				close(stopChan)
-				return
-			}
-		}
+		go server.New(config).Run(ctx)
 
+		<-signalChan
+		logger.Info("shutdown signal received, exiting...")
+		cancel()
 	},
 }
 

code-generation/pkg/codegen/assets/base.go.templ

@@ -1,6 +1,7 @@
 package base
 
 import (
+  "context"
   "github.com/awslabs/aws-service-operator/pkg/config"
   {{- range $index, $element := .Items}}
   "github.com/awslabs/aws-service-operator/pkg/operators/{{$element.Spec.Resource.Name}}"
@@ -19,16 +20,13 @@ func New(
   }
 }
 
-func (b *base) Watch(namespace string, stopCh chan struct{}) (err error) {
+func (b *base) Watch(ctx context.Context, namespace string) {
   {{- range $index, $element := .Items}}
   if b.config.Resources["{{$element.Spec.Resource.Name}}"] {
     {{$element.Spec.Resource.Name}}operator := {{$element.Spec.Resource.Name}}.NewOperator(b.config)
-  	err = {{$element.Spec.Resource.Name}}operator.StartWatch(namespace, stopCh)
-  	if err != nil {
-  	  return err
-  	}
+  	go {{$element.Spec.Resource.Name}}operator.StartWatch(ctx, namespace)
   }
 	{{- end}}
 
-	return nil
+	<-ctx.Done()
 }

code-generation/pkg/codegen/assets/operator.go.templ

@@ -6,6 +6,7 @@
 package {{.Spec.Resource.Name}}
 
 import (
+	"context"
 	{{- if .Spec.IsCustomized}}
 	 metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"github.com/awslabs/aws-service-operator/pkg/helpers"
@@ -43,7 +44,7 @@ func NewOperator(config *config.Config) *Operator {
 }
 
 // StartWatch watches for instances of Object Store custom resources and acts on them
-func (c *Operator) StartWatch(namespace string, stopCh chan struct{}) error {
+func (c *Operator) StartWatch(ctx context.Context, namespace string) {
 	resourceHandlers := cache.ResourceEventHandlerFuncs{
 		AddFunc:    c.onAdd,
 		UpdateFunc: c.onUpdate,
@@ -52,13 +53,11 @@ func (c *Operator) StartWatch(namespace string, stopCh chan struct{}) error {
 	{{- if .Spec.Queue}}
 	queuectrl := queue.New(c.config, c.config.AWSClientset, 1)
 	c.topicARN, _, _, _ = queuectrl.Register("{{.Spec.Resource.Name}}", &awsV1alpha1.{{.Spec.Kind}}{})
-	go queuectrl.StartWatch(queue.HandlerFunc(QueueUpdater), stopCh)
+	go queuectrl.StartWatch(queue.HandlerFunc(QueueUpdater), ctx.Done())
 	{{- end}}
 
 	oper := operator.New("{{.Spec.Resource.Plural}}", namespace, resourceHandlers, c.config.AWSClientset.RESTClient())
-	go oper.Watch(&awsV1alpha1.{{.Spec.Kind}}{}, stopCh)
-
-	return nil
+	oper.Watch(&awsV1alpha1.{{.Spec.Kind}}{}, ctx.Done())
 }
 
 {{- if .Spec.Queue}}

code-generation/pkg/codegen/templates.go

@@ -1,19 +1,19 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 
 	awsscheme "github.com/awslabs/aws-service-operator/pkg/client/clientset/versioned/scheme"
 	"github.com/awslabs/aws-service-operator/pkg/config"
 	opBase "github.com/awslabs/aws-service-operator/pkg/operators/base"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/sirupsen/logrus"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes/scheme"
 	typedcorev1 "k8s.io/client-go/kubernetes/typed/core/v1"
 	"k8s.io/client-go/tools/record"
-
-	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
 const controllerName = "aws-service-operator"
@@ -25,16 +25,27 @@ func New(config *config.Config) *Server {
 	}
 }
 
-func (c *Server) exposeMetrics(errChan chan error) {
-	http.Handle("/metrics", promhttp.Handler())
-	err := http.ListenAndServe(":9090", nil)
-	if err != nil {
-		errChan <- fmt.Errorf("unable to expose metrics: %v", err)
+func (c *Server) exposeMetrics(errChan chan error, ctx context.Context) {
+	c.Handle("/metrics", promhttp.Handler())
+	server := http.Server{
+		Addr:    ":9090",
+		Handler: c,
 	}
-}
+	defer server.Shutdown(ctx)
 
-func (c *Server) watchOperatorResources(errChan chan error, stopChan chan struct{}) {
+	go func() {
+		err := server.ListenAndServe()
+		if err != nil {
+			errChan <- fmt.Errorf("unable to expose metrics: %v", err)
+		}
+	}()
 
+	c.Config.Logger.Info("metrics server started")
+	<-ctx.Done()
+	c.Config.Logger.Info("metrics server stopped")
+}
+
+func (c *Server) watchOperatorResources(errChan chan error, ctx context.Context) {
 	logger := c.Config.Logger
 
 	logger.Info("getting kubernetes context")
@@ -48,27 +59,31 @@ func (c *Server) watchOperatorResources(errChan chan error, stopChan chan struct
 	// start watching the aws operator resources
 	logger.WithFields(logrus.Fields{"resources": c.Config.Resources}).Info("Watching")
 	operators := opBase.New(c.Config) // TODO: remove context and Clientset
-	err := operators.Watch(corev1.NamespaceAll, stopChan)
-	if err != nil {
-		errChan <- fmt.Errorf("unable to watch resources: %v", err)
-	}
+
+	go operators.Watch(ctx, corev1.NamespaceAll)
+	<-ctx.Done()
+	c.Config.Logger.Info("operators stopped")
 }
 
 // Run starts the server to listen to Kubernetes
-func (c *Server) Run(stopChan chan struct{}) {
+func (c *Server) Run(ctx context.Context) {
 	config := c.Config
 	logger := config.Logger
 	errChan := make(chan error, 1)
 
 	logger.Info("starting metrics server")
-	go c.exposeMetrics(errChan)
+	go c.exposeMetrics(errChan, ctx)
 
 	logger.Info("starting resource watcher")
-	go c.watchOperatorResources(errChan, stopChan)
+	go c.watchOperatorResources(errChan, ctx)
 
-	err := <-errChan
-	if err != nil {
-		logger.Fatal(err)
+	for {
+		select {
+		case err := <-errChan:
+			c.Config.Logger.WithError(err).Fatal(err)
+		case <-ctx.Done():
+			c.Config.Logger.Info("stop signal received. waiting for operators to stop")
+			return
+		}
 	}
-
 }

pkg/server/server.go

@@ -1,10 +1,13 @@
 package server
 
 import (
+	"net/http"
+
 	"github.com/awslabs/aws-service-operator/pkg/config"
 )
 
 // Server defines the bas construct for the operator
 type Server struct {
+	http.ServeMux
 	Config *config.Config
 }