Merge pull request #132 from tantona/fix-operator-shutdown

Fixed operator signal interruption using context WithCancel

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}}