From 6a063500e628f98bbdfc00e04bfc2d6861e15353 Mon Sep 17 00:00:00 2001
From: zwtop <wang.zhan@smartx.com>
Date: Thu, 30 May 2024 09:17:59 +0800
Subject: [PATCH] feat: graceful shutdown when leader lost

Signed-off-by: zwtop <wang.zhan@smartx.com>
---
 pkg/server/shutdown.go      | 42 ++++++++++++++++++++++
 pkg/server/shutdown_test.go | 69 +++++++++++++++++++++++++++++++++++++
 2 files changed, 111 insertions(+)
 create mode 100644 pkg/server/shutdown.go
 create mode 100644 pkg/server/shutdown_test.go

diff --git a/pkg/server/shutdown.go b/pkg/server/shutdown.go
new file mode 100644
index 0000000..92ce581
--- /dev/null
+++ b/pkg/server/shutdown.go
@@ -0,0 +1,42 @@
+package server
+
+import (
+	"context"
+
+	"k8s.io/klog/v2"
+
+	"github.com/everoute/runtime/pkg/options"
+)
+
+// GracefulShutdown graceful shutdown apiserver when leader lost
+// make sure api should available before the apiserver shutdown
+// stopped until context done, or another node becomes leader ready
+func GracefulShutdown(ctx context.Context, config *options.RecommendedConfig) {
+	electionCh := make(chan string)
+
+	go func() {
+		defer close(electionCh)
+		electionClient := config.LeaderElectionClient
+		if electionClient == nil {
+			return
+		}
+		// check if leader has been changed
+		if ld := electionClient.GetLeader(); ld != "" && ld != electionClient.Identity() {
+			electionCh <- ld
+			return
+		}
+		// until leading state update
+		for electionClient.UntilLeadingStateUpdate(ctx.Done()) {
+			if ld := electionClient.GetLeader(); ld != "" && ld != electionClient.Identity() {
+				electionCh <- ld
+			}
+		}
+	}()
+
+	select {
+	case <-ctx.Done():
+		klog.Fatalf("stopped when context done: %s", ctx.Err())
+	case ld := <-electionCh:
+		klog.Fatalf("stopped when node %s becomes leader", ld)
+	}
+}
diff --git a/pkg/server/shutdown_test.go b/pkg/server/shutdown_test.go
new file mode 100644
index 0000000..6719306
--- /dev/null
+++ b/pkg/server/shutdown_test.go
@@ -0,0 +1,69 @@
+package server_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/agiledragon/gomonkey/v2"
+	. "github.com/onsi/gomega"
+	"k8s.io/apimachinery/pkg/util/rand"
+	"k8s.io/klog/v2"
+
+	"github.com/everoute/runtime/pkg/options"
+	"github.com/everoute/runtime/pkg/server"
+	. "github.com/everoute/runtime/pkg/util/testing"
+)
+
+func TestGracefulShutdown(t *testing.T) {
+	RegisterTestingT(t)
+
+	ctx := context.Background()
+	messageCh := make(chan string, 10)
+	patches := gomonkey.ApplyFunc(klog.Fatalf, func(format string, args ...interface{}) {
+		klog.Infof(format, args...)
+		messageCh <- fmt.Sprintf(format, args...)
+	})
+	defer patches.Reset()
+
+	t.Run("should shutdown when another node has been becomes leader", func(t *testing.T) {
+		electionClient := NewFakeLeaderElectionClient(rand.String(20))
+		electionClient.SetLeader(rand.String(20))
+		go server.GracefulShutdown(ctx, &options.RecommendedConfig{LeaderElectionClient: electionClient})
+
+		select {
+		case msg := <-messageCh:
+			Expect(msg).Should(ContainSubstring("becomes leader"))
+		case <-time.After(time.Second):
+			t.Fatalf("unexpect timeout wait graceful shutdown")
+		}
+	})
+
+	t.Run("should shutdown when another node becomes leader", func(t *testing.T) {
+		electionClient := NewFakeLeaderElectionClient(rand.String(20))
+		go func() { time.Sleep(200 * time.Millisecond); electionClient.SetLeader(rand.String(20)) }()
+		go server.GracefulShutdown(ctx, &options.RecommendedConfig{LeaderElectionClient: electionClient})
+
+		select {
+		case msg := <-messageCh:
+			Expect(msg).Should(ContainSubstring("becomes leader"))
+		case <-time.After(time.Second):
+			t.Fatalf("unexpect timeout wait graceful shutdown")
+		}
+	})
+
+	t.Run("should shutdown when context done", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(ctx, 200*time.Millisecond)
+		defer cancel()
+		electionClient := NewFakeLeaderElectionClient(rand.String(20))
+		go server.GracefulShutdown(ctx, &options.RecommendedConfig{LeaderElectionClient: electionClient})
+
+		select {
+		case msg := <-messageCh:
+			Expect(msg).Should(ContainSubstring("stopped when context done"))
+		case <-time.After(time.Second):
+			t.Fatalf("unexpect timeout wait graceful shutdown")
+		}
+	})
+}