From 83ae095ab9197f168a6bd3f6bd355f89bce39a9c Mon Sep 17 00:00:00 2001
From: Brad Davidson <brad.davidson@rancher.com>
Date: Tue, 18 Jun 2024 23:30:28 +0000
Subject: [PATCH] Replace 1-weight semaphore on snapshots with simple mutex

Fixes an issue where the semaphore wasn't permanently initialized
until a scheduled snapshot was taken, allowing multiple on-demand
snapshots to be taken until the first scheduled snapshot was triggered.

Signed-off-by: Brad Davidson <brad.davidson@rancher.com>
---
 pkg/etcd/etcd.go             | 23 ++++++++++++-----------
 pkg/etcd/snapshot.go         | 28 ++++++----------------------
 pkg/etcd/snapshot_handler.go | 10 +++++-----
 3 files changed, 23 insertions(+), 38 deletions(-)

diff --git a/pkg/etcd/etcd.go b/pkg/etcd/etcd.go
index 8d087f227a11..e3907e696e1f 100644
--- a/pkg/etcd/etcd.go
+++ b/pkg/etcd/etcd.go
@@ -16,6 +16,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
@@ -43,7 +44,6 @@ import (
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.etcd.io/etcd/etcdutl/v3/snapshot"
 	"go.uber.org/zap"
-	"golang.org/x/sync/semaphore"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
@@ -105,14 +105,14 @@ var _ managed.Driver = &ETCD{}
 type MemberStatus string
 
 type ETCD struct {
-	client      *clientv3.Client
-	config      *config.Control
-	name        string
-	address     string
-	cron        *cron.Cron
-	s3          *S3
-	cancel      context.CancelFunc
-	snapshotSem *semaphore.Weighted
+	client     *clientv3.Client
+	config     *config.Control
+	name       string
+	address    string
+	cron       *cron.Cron
+	s3         *S3
+	cancel     context.CancelFunc
+	snapshotMu *sync.Mutex
 }
 
 type learnerProgress struct {
@@ -166,10 +166,11 @@ func (e *MemberListError) Is(target error) bool {
 func errMemberListFailed() error { return &MemberListError{} }
 
 // NewETCD creates a new value of type
-// ETCD with an initialized cron value.
+// ETCD with initialized cron and snapshot mutex values.
 func NewETCD() *ETCD {
 	return &ETCD{
-		cron: cron.New(cron.WithLogger(cronLogger)),
+		cron:       cron.New(cron.WithLogger(cronLogger)),
+		snapshotMu: &sync.Mutex{},
 	}
 }
 
diff --git a/pkg/etcd/snapshot.go b/pkg/etcd/snapshot.go
index 664352eab053..8669b8443ffa 100644
--- a/pkg/etcd/snapshot.go
+++ b/pkg/etcd/snapshot.go
@@ -29,7 +29,6 @@ import (
 	"github.com/robfig/cron/v3"
 	"github.com/sirupsen/logrus"
 	"go.etcd.io/etcd/etcdutl/v3/snapshot"
-	"golang.org/x/sync/semaphore"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/equality"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
@@ -44,10 +43,9 @@ import (
 )
 
 const (
-	maxConcurrentSnapshots = 1
-	compressedExtension    = ".zip"
-	metadataDir            = ".metadata"
-	errorTTL               = 24 * time.Hour
+	compressedExtension = ".zip"
+	metadataDir         = ".metadata"
+	errorTTL            = 24 * time.Hour
 )
 
 var (
@@ -106,16 +104,6 @@ func snapshotDir(config *config.Control, create bool) (string, error) {
 	return snapshotDir, nil
 }
 
-// preSnapshotSetup checks to see if the necessary components are in place
-// to perform an Etcd snapshot. This is necessary primarily for on-demand
-// snapshots since they're performed before normal Etcd setup is completed.
-func (e *ETCD) preSnapshotSetup(ctx context.Context) error {
-	if e.snapshotSem == nil {
-		e.snapshotSem = semaphore.NewWeighted(maxConcurrentSnapshots)
-	}
-	return nil
-}
-
 // compressSnapshot compresses the given snapshot and provides the
 // caller with the path to the file.
 func (e *ETCD) compressSnapshot(snapshotDir, snapshotName, snapshotPath string, now time.Time) (string, error) {
@@ -208,14 +196,10 @@ func (e *ETCD) decompressSnapshot(snapshotDir, snapshotFile string) (string, err
 // subcommand for prune that can be run manually if the user wants to remove old snapshots.
 // Returns metadata about the new and pruned snapshots.
 func (e *ETCD) Snapshot(ctx context.Context) (*managed.SnapshotResult, error) {
-	if err := e.preSnapshotSetup(ctx); err != nil {
-		return nil, err
-	}
-	if !e.snapshotSem.TryAcquire(maxConcurrentSnapshots) {
-		return nil, fmt.Errorf("%d snapshots already in progress", maxConcurrentSnapshots)
+	if !e.snapshotMu.TryLock() {
+		return nil, errors.New("snapshot save already in progress")
 	}
-	defer e.snapshotSem.Release(maxConcurrentSnapshots)
-
+	defer e.snapshotMu.Unlock()
 	// make sure the core.Factory is initialized before attempting to add snapshot metadata
 	var extraMetadata *v1.ConfigMap
 	if e.config.Runtime.Core == nil {
diff --git a/pkg/etcd/snapshot_handler.go b/pkg/etcd/snapshot_handler.go
index a7f336454843..0bae2e0401b7 100644
--- a/pkg/etcd/snapshot_handler.go
+++ b/pkg/etcd/snapshot_handler.go
@@ -150,11 +150,11 @@ func (e *ETCD) withRequest(sr *SnapshotRequest) *ETCD {
 			EtcdSnapshotName:      e.config.EtcdSnapshotName,
 			EtcdSnapshotRetention: e.config.EtcdSnapshotRetention,
 		},
-		name:        e.name,
-		address:     e.address,
-		cron:        e.cron,
-		cancel:      e.cancel,
-		snapshotSem: e.snapshotSem,
+		name:       e.name,
+		address:    e.address,
+		cron:       e.cron,
+		cancel:     e.cancel,
+		snapshotMu: e.snapshotMu,
 	}
 	if len(sr.Name) > 0 {
 		re.config.EtcdSnapshotName = sr.Name[0]