cmd/run: wait for address to be reported before removing taint

cmd/main.go

@@ -137,6 +137,63 @@ func assignAddress(c context.Context, log *logrus.Entry, client kubernetes.Inter
 	return "", errors.New("reached maximum number of retries")
 }
 
+func waitForAddressToBeReported(c context.Context, log *logrus.Entry, explorer nd.Explorer, node *types.Node, address string, cfg *config.Config) error {
+	ctx, cancel := context.WithCancel(c)
+	defer cancel()
+
+	// ticker for retry interval
+	ticker := time.NewTicker(cfg.RetryInterval)
+	defer ticker.Stop()
+
+	for retryCounter := 0; retryCounter <= cfg.RetryAttempts; retryCounter++ {
+		log.WithFields(logrus.Fields{
+			"node":           node.Name,
+			"instance":       node.Instance,
+			"address":        address,
+			"retry-counter":  retryCounter,
+			"retry-attempts": cfg.RetryAttempts,
+		}).Debug("Waiting for node to report assigned address")
+
+		node, err := explorer.GetNode(ctx, node.Name)
+		if err == nil {
+			for _, ip := range node.ExternalIPs {
+				if ip.String() == address {
+					log.WithFields(logrus.Fields{
+						"node":           node.Name,
+						"instance":       node.Instance,
+						"address":        address,
+						"retry-counter":  retryCounter,
+						"retry-attempts": cfg.RetryAttempts,
+					}).Info("Node is reporting assigned address")
+					return nil
+				}
+			}
+			log.WithError(err).WithFields(logrus.Fields{
+				"node":     node.Name,
+				"instance": node.Instance,
+				"address":  address,
+			}).Error("Node is not yet reporting the assigned address")
+		} else {
+			log.WithError(err).WithFields(logrus.Fields{
+				"node":     node.Name,
+				"instance": node.Instance,
+				"address":  address,
+			}).Error("failed to check if node is reporting the assigned address")
+		}
+
+		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 waiting for node to report assigned address")
+		}
+	}
+	return errors.New("reached maximum number of retries")
+}
+
 func run(c context.Context, log *logrus.Entry, cfg *config.Config) error {
 	ctx, cancel := context.WithCancel(c)
 	defer cancel()
@@ -170,12 +227,16 @@ func run(c context.Context, log *logrus.Entry, cfg *config.Config) error {
 		return errors.Wrap(err, "initializing assigner")
 	}
 
-	_, err = assignAddress(ctx, log, clientset, assigner, n, cfg)
+	assignedAddress, err := assignAddress(ctx, log, clientset, assigner, n, cfg)
 	if err != nil {
 		return errors.Wrap(err, "assigning static public IP address")
 	}
 
 	if cfg.TaintKey != "" {
+		if err := waitForAddressToBeReported(ctx, log, explorer, n, assignedAddress, cfg); err != nil {
+			return errors.Wrap(err, "waiting for node to report assigned address")
+		}
+
 		logger := log.WithField("taint-key", cfg.TaintKey)
 		tainter := nd.NewTainter(clientset)