ce7cb702218693d08ea6d83464d2457c9250027c

services/short/tunnel/lease.go

@@ -0,0 +1,92 @@
+package tunnel
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/tools/leaderelection"
+	"k8s.io/client-go/tools/leaderelection/resourcelock"
+)
+
+func label(ctx context.Context, kube kubernetes.Interface, namespace, name string, value any) error {
+	data, err := json.Marshal(map[string]any{"metadata": map[string]any{"labels": map[string]any{"role": value}}})
+	if err != nil {
+		return err
+	}
+	_, err = kube.CoreV1().Pods(namespace).Patch(ctx, name, types.MergePatchType, data, metav1.PatchOptions{})
+	return err
+}
+
+func claim(ctx context.Context, kube kubernetes.Interface, namespace, name string) error {
+	pods, err := kube.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{LabelSelector: "app.kubernetes.io/name=tunnel,app.kubernetes.io/instance=short"})
+	if err != nil {
+		return err
+	}
+	found := false
+	for _, pod := range pods.Items {
+		if pod.Name == name {
+			found = true
+		} else if pod.Labels["role"] == "active" {
+			if err := label(ctx, kube, namespace, pod.Name, nil); err != nil {
+				return err
+			}
+		}
+	}
+	if !found {
+		return fmt.Errorf("tunnel pod is missing")
+	}
+	return label(ctx, kube, namespace, name, "active")
+}
+
+func ownership(ctx context.Context, kube kubernetes.Interface, namespace, name string, work func(context.Context) error) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	var gate sync.Mutex
+	var stopped bool
+	var running sync.WaitGroup
+	result := make(chan error, 1)
+	lock := &resourcelock.LeaseLock{LeaseMeta: metav1.ObjectMeta{Name: "tunnel", Namespace: namespace}, Client: kube.CoordinationV1(), LockConfig: resourcelock.ResourceLockConfig{Identity: name}}
+	elector, err := leaderelection.NewLeaderElector(leaderelection.LeaderElectionConfig{
+		// A successful renewal reply can consume 20s before retry, renewal failure, and child shutdown.
+		Lock: lock, LeaseDuration: 60 * time.Second, RenewDeadline: 20 * time.Second, RetryPeriod: 5 * time.Second,
+		// Immediate release could start a replacement while charon is stopping.
+		ReleaseOnCancel: false,
+		Callbacks: leaderelection.LeaderCallbacks{
+			OnStartedLeading: func(leader context.Context) {
+				gate.Lock()
+				if stopped {
+					gate.Unlock()
+					return
+				}
+				running.Add(1)
+				gate.Unlock()
+				defer running.Done()
+				result <- work(leader)
+				cancel()
+			},
+			OnStoppedLeading: func() {
+				gate.Lock()
+				stopped = true
+				cancel()
+				gate.Unlock()
+				running.Wait()
+			},
+		},
+	})
+	if err != nil {
+		return err
+	}
+	elector.Run(ctx)
+	select {
+	case err := <-result:
+		return err
+	default:
+		return nil
+	}
+}