From ac92e06b9cc169f2b9632b4b07846c7471f43138 Mon Sep 17 00:00:00 2001
From: Gabe Cook <gabe565@gmail.com>
Date: Sat, 7 Oct 2023 02:09:45 -0500
Subject: [PATCH] perf: Improve startup time by watching for job pod

---
 internal/util/cmd_setup.go | 46 ++++++++++++++++++++++++++++++++------
 1 file changed, 39 insertions(+), 7 deletions(-)

diff --git a/internal/util/cmd_setup.go b/internal/util/cmd_setup.go
index 022c8d55..63880830 100644
--- a/internal/util/cmd_setup.go
+++ b/internal/util/cmd_setup.go
@@ -164,7 +164,7 @@ func DefaultSetup(cmd *cobra.Command, conf *config.Global, opts SetupOptions) (e
 			return err
 		}
 
-		if err := waitForPod(cmd, conf); err != nil {
+		if err := watchJobPod(cmd, conf); err != nil {
 			Teardown(cmd, conf)
 			return err
 		}
@@ -278,7 +278,7 @@ func createJob(cmd *cobra.Command, conf *config.Global, actionName string) error
 	return nil
 }
 
-func waitForPod(cmd *cobra.Command, conf *config.Global) error {
+func watchJobPod(cmd *cobra.Command, conf *config.Global) error {
 	log.WithFields(log.Fields{
 		"namespace": conf.Namespace,
 		"name":      "job.batch/" + conf.Job.ObjectMeta.Name,
@@ -287,13 +287,45 @@ func waitForPod(cmd *cobra.Command, conf *config.Global) error {
 	ctx, cancel := context.WithTimeout(cmd.Context(), 5*time.Minute)
 	defer cancel()
 
+	watch, err := conf.Client.ClientSet.CoreV1().Pods(conf.Namespace).Watch(ctx, metav1.ListOptions{
+		LabelSelector: "batch.kubernetes.io/controller-uid=" + string(conf.Job.ObjectMeta.UID),
+	})
+	if err != nil {
+		return pollJobPod(ctx, conf)
+	}
+	defer func() {
+		watch.Stop()
+	}()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case event := <-watch.ResultChan():
+			if pod, ok := event.Object.(*corev1.Pod); ok {
+				switch pod.Status.Phase {
+				case corev1.PodRunning:
+					conf.Host = conf.DbPod.Status.PodIP
+					pod.DeepCopyInto(&conf.JobPod)
+					return nil
+				case corev1.PodFailed:
+					return errors.New("pod failed")
+				case corev1.PodSucceeded:
+					return errors.New("pod exited early")
+				}
+			} else {
+				return errors.New("unexpected runtime object type")
+			}
+		}
+	}
+}
+
+func pollJobPod(ctx context.Context, conf *config.Global) error {
 	return wait.PollUntilContextCancel(
 		ctx, time.Second, true, func(ctx context.Context) (done bool, err error) {
-			list, err := conf.Client.ClientSet.CoreV1().Pods(conf.Namespace).List(
-				cmd.Context(), metav1.ListOptions{
-					LabelSelector: "batch.kubernetes.io/controller-uid=" + string(conf.Job.ObjectMeta.UID),
-				},
-			)
+			list, err := conf.Client.ClientSet.CoreV1().Pods(conf.Namespace).List(ctx, metav1.ListOptions{
+				LabelSelector: "batch.kubernetes.io/controller-uid=" + string(conf.Job.ObjectMeta.UID),
+			})
 			if err != nil {
 				return false, err
 			}