add leader election support for external snapshotter

Signed-off-by: Andrew Sy Kim <[email protected]>

Author: andrewsykim

cmd/csi-snapshotter/main.go

@@ -30,6 +30,7 @@ import (
 	"k8s.io/client-go/tools/clientcmd"
 	"k8s.io/klog"
 
+	"github.com/kubernetes-csi/csi-lib-utils/leaderelection"
 	"github.com/kubernetes-csi/external-snapshotter/pkg/connection"
 	"github.com/kubernetes-csi/external-snapshotter/pkg/controller"
 
@@ -60,10 +61,14 @@ var (
 	snapshotNamePrefix              = flag.String("snapshot-name-prefix", "snapshot", "Prefix to apply to the name of a created snapshot")
 	snapshotNameUUIDLength          = flag.Int("snapshot-name-uuid-length", -1, "Length in characters for the generated uuid of a created snapshot. Defaults behavior is to NOT truncate.")
 	showVersion                     = flag.Bool("version", false, "Show version.")
+
+	leaderElection          = flag.Bool("leader-election", false, "enables leader election")
+	leaderElectionNamespace = flag.String("leader-election-namespace", "", "the namespace where the leader election lock exists, should be the same namespace attacher runs")
 )
 
 var (
-	version = "unknown"
+	version                = "unknown"
+	leaderElectionLockName = "external-snapshotter-leader-election"
 )
 
 func main() {
@@ -183,17 +188,33 @@ func main() {
 		*snapshotNameUUIDLength,
 	)
 
-	// run...
-	stopCh := make(chan struct{})
-	factory.Start(stopCh)
-	coreFactory.Start(stopCh)
-	go ctrl.Run(threads, stopCh)
-
-	// ...until SIGINT
-	c := make(chan os.Signal, 1)
-	signal.Notify(c, os.Interrupt)
-	<-c
-	close(stopCh)
+	run := func(context.Context) {
+		// run...
+		stopCh := make(chan struct{})
+		factory.Start(stopCh)
+		coreFactory.Start(stopCh)
+		go ctrl.Run(threads, stopCh)
+
+		// ...until SIGINT
+		c := make(chan os.Signal, 1)
+		signal.Notify(c, os.Interrupt)
+		<-c
+		close(stopCh)
+	}
+
+	if !*leaderElection {
+		run(context.TODO())
+	} else {
+		if *leaderElectionNamespace == "" {
+			klog.Error("--leader-election-namespace is required")
+			os.Exit(1)
+		}
+
+		le := leaderelection.NewLeaderElection(kubeClient, leaderElectionLockName, *leaderElectionNamespace, run)
+		if err := le.Run(); err != nil {
+			klog.Fatalf("failed to initialze leader election: %v", err)
+		}
+	}
 }
 
 func buildConfig(kubeconfig string) (*rest.Config, error) {