From 2357a626b199a55760e264ec737c58d05c94f21e Mon Sep 17 00:00:00 2001
From: Vadim Berezniker <vadim@berezniker.com>
Date: Wed, 25 Sep 2024 17:59:52 -0700
Subject: [PATCH] Wait for pod evictions when draining, up to 60 seconds.

---
 internal/drainer/drain.go | 45 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/internal/drainer/drain.go b/internal/drainer/drain.go
index 3860529..54cd774 100644
--- a/internal/drainer/drain.go
+++ b/internal/drainer/drain.go
@@ -3,6 +3,8 @@ package drain
 import (
 	"context"
 	"fmt"
+	"strings"
+	"time"
 
 	"github.com/sirupsen/logrus"
 	"k8s.io/api/core/v1"
@@ -13,6 +15,11 @@ import (
 	"k8s.io/client-go/kubernetes"
 )
 
+const (
+	// Wait up to this amount of time for pods to be evicted from a node.
+	drainMaxWait = 60 * time.Second
+)
+
 // Config configures a Drainer.
 type Config struct {
 	Client kubernetes.Interface
@@ -74,6 +81,7 @@ func (d *drainer) Drain(ctx context.Context, node string) error {
 		return err
 	}
 
+	evictedPods := []string{}
 	for _, pod := range pods {
 		fields["pod"] = pod.GetName()
 		d.log.WithFields(fields).Info("drainer: evicting pod")
@@ -83,6 +91,43 @@ func (d *drainer) Drain(ctx context.Context, node string) error {
 			d.log.WithFields(fields).Errorf("drainer: error evicting pod: %v", err)
 			return err
 		}
+		evictedPods = append(evictedPods, pod.GetName())
+	}
+
+	start := time.Now()
+	for len(evictedPods) > 0 {
+		podsDescription := ""
+		if len(evictedPods) > 5 {
+			podsDescription = strings.Join(evictedPods[:5], ", ") + " ..."
+		} else {
+			podsDescription = strings.Join(evictedPods, ", ")
+		}
+		d.log.WithFields(fields).Infof("drainer: waiting for %d pods to be evicted: %s", len(evictedPods), podsDescription)
+
+		if time.Since(start) > drainMaxWait {
+			d.log.WithFields(fields).Infof("drainer: waited maximum amount of time for evictions, continuing")
+			break
+		}
+
+		pods, err := d.getPodsForDeletion(ctx, node)
+		if err != nil {
+			d.log.WithFields(fields).Errorf("drainer: error getting pods: %v", err)
+			return err
+		}
+		podsByName := make(map[string]struct{}, len(pods))
+		for _, pod := range pods {
+			podsByName[pod.GetName()] = struct{}{}
+		}
+		remainingPods := []string{}
+		for _, pod := range evictedPods {
+			if _, ok := podsByName[pod]; ok {
+				remainingPods = append(remainingPods, pod)
+			}
+		}
+		evictedPods = remainingPods
+		if len(evictedPods) > 0 {
+			time.Sleep(1 * time.Second)
+		}
 	}
 
 	d.log.WithFields(fields).Info("drainer: drained node")