improve signal termination cleanup (#137)

Author: alexei-led

cmd/main.go

@@ -23,7 +23,8 @@ import (
 type contextKey string
 
 const (
-	developModeKey contextKey = "develop-mode"
+	developModeKey  contextKey = "develop-mode"
+	unassignTimeout            = 5 * time.Minute
 )
 
 var (
@@ -81,36 +82,29 @@ func prepareLogger(level string, json bool) *logrus.Entry {
 func assignAddress(c context.Context, log *logrus.Entry, assigner address.Assigner, node *types.Node, cfg *config.Config) error {
 	ctx, cancel := context.WithCancel(c)
 	defer cancel()
-	// retry counter
-	retryCounter := 0
+
 	// ticker for retry interval
 	ticker := time.NewTicker(cfg.RetryInterval)
 	defer ticker.Stop()
 
-	for {
+	for retryCounter := 0; retryCounter <= cfg.RetryAttempts; retryCounter++ {
 		err := assigner.Assign(ctx, node.Instance, node.Zone, cfg.Filter, cfg.OrderBy)
-		if err != nil && errors.Is(err, address.ErrStaticIPAlreadyAssigned) {
-			log.Infof("static public IP address already assigned to node instance %s", node.Instance)
+		if err == nil || errors.Is(err, address.ErrStaticIPAlreadyAssigned) {
 			return nil
 		}
-		if err != nil {
-			log.WithError(err).Errorf("failed to assign static public IP address to node %s", node.Name)
-			if retryCounter < cfg.RetryAttempts {
-				retryCounter++
-				log.Infof("retrying after %v", cfg.RetryInterval)
-			} else {
-				log.Infof("reached maximum number of retries (%d)", cfg.RetryAttempts)
-				return errors.Wrap(err, "reached maximum number of retries")
-			}
-			select {
-			case <-ticker.C:
-				continue
-			case <-ctx.Done():
-				return errors.Wrap(err, "context is done")
-			}
+
+		log.WithError(err).Errorf("failed to assign static public IP address to node %s", node.Name)
+		log.Infof("retrying after %v", cfg.RetryInterval)
+
+		select {
+		case <-ticker.C:
+			continue
+		case <-ctx.Done():
+			// If the context is done, return an error indicating that the operation was cancelled
+			return errors.Wrap(ctx.Err(), "context cancelled while assigning addresses")
 		}
-		return nil
 	}
+	return errors.New("reached maximum number of retries")
 }
 
 func run(c context.Context, log *logrus.Entry, cfg *config.Config) error {
@@ -145,8 +139,9 @@ func run(c context.Context, log *logrus.Entry, cfg *config.Config) error {
 		return errors.Wrap(err, "initializing assigner")
 	}
 	// assign static public IP address
-	errorCh := make(chan error)
+	errorCh := make(chan error, 1) // buffered channel to avoid goroutine leak
 	go func() {
+		defer close(errorCh) // close the channel when the goroutine exits to avoid goroutine leak
 		e := assignAddress(ctx, log, assigner, n, cfg)
 		if e != nil {
 			errorCh <- e
@@ -159,13 +154,16 @@ func run(c context.Context, log *logrus.Entry, cfg *config.Config) error {
 			return errors.Wrap(err, "assigning static public IP address")
 		}
 	case <-ctx.Done():
-		log.Infof("kubeip agent stopped")
+		log.Infof("kubeip agent gracefully stopped")
 		if cfg.ReleaseOnExit {
 			log.Infof("releasing static public IP address")
+			releaseCtx, releaseCancel := context.WithTimeout(context.Background(), unassignTimeout) // release the static public IP address within 5 minutes
+			defer releaseCancel()
 			// use a different context for releasing the static public IP address since the main context is canceled
-			if err = assigner.Unassign(context.Background(), n.Instance, n.Zone); err != nil { //nolint:contextcheck
-				return errors.Wrap(err, "releasing static public IP address")
+			if err = assigner.Unassign(releaseCtx, n.Instance, n.Zone); err != nil { //nolint:contextcheck
+				return errors.Wrap(err, "failed to release static public IP address")
 			}
+			log.Infof("static public IP address released")
 		}
 	}
 

cmd/main_test.go

@@ -97,6 +97,38 @@ func Test_assignAddress(t *testing.T) {
 			},
 			wantErr: true,
 		},
+		{
+			name: "context cancelled while assigning addresses",
+			args: args{
+				c: func() context.Context {
+					ctx, cancel := context.WithCancel(context.Background())
+					go func() {
+						// Simulate a shutdown signal being received after a short delay
+						time.Sleep(20 * time.Millisecond)
+						cancel()
+					}()
+					return ctx
+				}(),
+				assignerFn: func(t *testing.T) address.Assigner {
+					mock := mocks.NewAssigner(t)
+					mock.EXPECT().Assign(tmock.Anything, "test-instance", "test-zone", []string{"test-filter"}, "test-order-by").Return(errors.New("error")).Maybe()
+					return mock
+				},
+				node: &types.Node{
+					Name:     "test-node",
+					Instance: "test-instance",
+					Region:   "test-region",
+					Zone:     "test-zone",
+				},
+				cfg: &config.Config{
+					Filter:        []string{"test-filter"},
+					OrderBy:       "test-order-by",
+					RetryAttempts: 10,
+					RetryInterval: 5 * time.Millisecond,
+				},
+			},
+			wantErr: true,
+		},
 		{
 			name: "error after a few retries and context is done",
 			args: args{

internal/address/gcp.go

@@ -246,8 +246,12 @@ func (a *gcpAssigner) Assign(ctx context.Context, instanceID, zone string, filte
 	// try to assign all available addresses until one succeeds
 	// due to concurrency, it is possible that another kubeip instance will assign the same address
 	for _, address := range addresses {
+		// check if context is done before trying to assign an address
+		if ctx.Err() != nil {
+			return errors.Wrap(ctx.Err(), "context cancelled while assigning addresses")
+		}
 		if err = tryAssignAddress(ctx, a, instance, a.region, zone, address); err != nil {
-			a.logger.WithError(err).Errorf("failed to assign static public IP address %s", address.Address)
+			a.logger.WithError(err).WithField("address", address.Address).Error("failed to assign static public IP address")
 			continue
 		}
 		// break the loop after successfully assigning an address
@@ -400,6 +404,10 @@ func (a *gcpAssigner) createUserMap(assigned []*compute.Address) map[string]stru
 
 func retryAddEphemeralAddress(ctx context.Context, logger *logrus.Entry, as internalAssigner, instance *compute.Instance, zone string) error {
 	for i := 0; i < maxRetries; i++ {
+		// check if context is done before trying to assign an address
+		if ctx.Err() != nil {
+			return errors.Wrap(ctx.Err(), "context cancelled while assigning ephemeral addresses")
+		}
 		if err := as.AddInstanceAddress(ctx, instance, zone, nil); err != nil {
 			logger.WithError(err).Error("failed to assign ephemeral public IP address, retrying")
 			continue