From 46b6f446fdd62ac095f1a3eaab3bc693b6b5e9a7 Mon Sep 17 00:00:00 2001
From: Jose Diaz-Gonzalez <email@josediazgonzalez.com>
Date: Thu, 18 Jan 2024 00:19:20 -0500
Subject: [PATCH] feat: implement ps:stop for k3s scheduler

This simply scales down processes.
---
 plugins/scheduler-k3s/Makefile                |  2 +-
 .../scheduler-k3s/src/triggers/triggers.go    |  4 ++
 plugins/scheduler-k3s/triggers.go             | 54 +++++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/plugins/scheduler-k3s/Makefile b/plugins/scheduler-k3s/Makefile
index 7efff411e..8609e8ee3 100644
--- a/plugins/scheduler-k3s/Makefile
+++ b/plugins/scheduler-k3s/Makefile
@@ -1,5 +1,5 @@
 SUBCOMMANDS = subcommands/initialize subcommands/report subcommands/set
-TRIGGERS = triggers/install triggers/post-app-clone-setup triggers/post-app-rename-setup triggers/post-delete triggers/post-registry-login triggers/report triggers/scheduler-deploy triggers/scheduler-enter triggers/scheduler-logs
+TRIGGERS = triggers/install triggers/post-app-clone-setup triggers/post-app-rename-setup triggers/post-delete triggers/post-registry-login triggers/report triggers/scheduler-deploy triggers/scheduler-enter triggers/scheduler-logs triggers/scheduler-stop
 BUILD = commands subcommands triggers
 PLUGIN_NAME = scheduler-k3s
 
diff --git a/plugins/scheduler-k3s/src/triggers/triggers.go b/plugins/scheduler-k3s/src/triggers/triggers.go
index 4ac250d8f..7b8e74f0f 100644
--- a/plugins/scheduler-k3s/src/triggers/triggers.go
+++ b/plugins/scheduler-k3s/src/triggers/triggers.go
@@ -82,6 +82,10 @@ func main() {
 		}
 
 		err = scheduler_k3s.TriggerSchedulerLogs(scheduler, appName, processType, tail, quiet, numLines)
+	case "scheduler-stop":
+		scheduler := flag.Arg(0)
+		appName := flag.Arg(1)
+		err = scheduler_k3s.TriggerSchedulerStop(scheduler, appName)
 	default:
 		err = fmt.Errorf("Invalid plugin trigger call: %s", trigger)
 	}
diff --git a/plugins/scheduler-k3s/triggers.go b/plugins/scheduler-k3s/triggers.go
index 76a4e3a5c..9f27fb55d 100644
--- a/plugins/scheduler-k3s/triggers.go
+++ b/plugins/scheduler-k3s/triggers.go
@@ -22,6 +22,7 @@ import (
 	orderedmap "github.com/wk8/go-ordered-map/v2"
 	"gopkg.in/yaml.v3"
 	appsv1 "k8s.io/api/apps/v1"
+	autoscalingv1 "k8s.io/api/autoscaling/v1"
 	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -638,3 +639,56 @@ func TriggerSchedulerLogs(scheduler string, appName string, processType string,
 
 	return nil
 }
+
+// TriggerSchedulerStop stops an application
+func TriggerSchedulerStop(scheduler string, appName string) error {
+	if scheduler != "k3s" {
+		return nil
+	}
+
+	clientset, err := NewKubernetesClient()
+	if err != nil {
+		return fmt.Errorf("Error creating kubernetes client: %w", err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, os.Interrupt, syscall.SIGHUP,
+		syscall.SIGINT,
+		syscall.SIGQUIT,
+		syscall.SIGTERM)
+	go func() {
+		<-signals
+		cancel()
+	}()
+
+	namespace := common.PropertyGetDefault("scheduler-k3s", appName, "namespace", "default")
+	deploymentList, err := clientset.Client.AppsV1().Deployments(namespace).List(ctx, metav1.ListOptions{
+		LabelSelector: fmt.Sprintf("dokku.com/app-name=%s", appName),
+	})
+	if err != nil {
+		return fmt.Errorf("Error listing deployments: %w", err)
+	}
+
+	for _, deployment := range deploymentList.Items {
+		processType, ok := deployment.Annotations["dokku.com/process-type"]
+		if !ok {
+			return fmt.Errorf("Deployment %s does not have a process type annotation", deployment.Name)
+		}
+		common.LogVerboseQuiet(fmt.Sprintf("Stopping %s process", processType))
+		_, err := clientset.Client.AppsV1().Deployments(namespace).UpdateScale(ctx, deployment.Name, &autoscalingv1.Scale{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      deployment.Name,
+				Namespace: namespace,
+			},
+			Spec: autoscalingv1.ScaleSpec{
+				Replicas: 0,
+			},
+		}, metav1.UpdateOptions{})
+		if err != nil {
+			return fmt.Errorf("Error updating deployment scale: %w", err)
+		}
+	}
+
+	return nil
+}