From 7ae3d25f9103c33045035c54ae2fc13aa35d3853 Mon Sep 17 00:00:00 2001
From: Piotr Tabor <ptab@google.com>
Date: Fri, 2 Apr 2021 18:01:26 +0200
Subject: [PATCH] Membership: Add additional methods to trim/manage membership
 data in backend.

---
 server/etcdserver/api/membership/member.go | 34 +++++++++++++---------
 server/etcdserver/api/membership/store.go  |  8 +++++
 server/mvcc/backend/batch_tx.go            | 13 +++++++++
 server/mvcc/kvstore_test.go                |  1 +
 4 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/server/etcdserver/api/membership/member.go b/server/etcdserver/api/membership/member.go
index 8ab1f4d41..efecf9acc 100644
--- a/server/etcdserver/api/membership/member.go
+++ b/server/etcdserver/api/membership/member.go
@@ -49,27 +49,22 @@ type Member struct {
 // NewMember creates a Member without an ID and generates one based on the
 // cluster name, peer URLs, and time. This is used for bootstrapping/adding new member.
 func NewMember(name string, peerURLs types.URLs, clusterName string, now *time.Time) *Member {
-	return newMember(name, peerURLs, clusterName, now, false)
+	memberId := computeMemberId(peerURLs, clusterName, now)
+	return newMember(name, peerURLs, memberId, false)
 }
 
 // NewMemberAsLearner creates a learner Member without an ID and generates one based on the
 // cluster name, peer URLs, and time. This is used for adding new learner member.
 func NewMemberAsLearner(name string, peerURLs types.URLs, clusterName string, now *time.Time) *Member {
-	return newMember(name, peerURLs, clusterName, now, true)
+	memberId := computeMemberId(peerURLs, clusterName, now)
+	return newMember(name, peerURLs, memberId, true)
 }
 
-func newMember(name string, peerURLs types.URLs, clusterName string, now *time.Time, isLearner bool) *Member {
-	m := &Member{
-		RaftAttributes: RaftAttributes{
-			PeerURLs:  peerURLs.StringSlice(),
-			IsLearner: isLearner,
-		},
-		Attributes: Attributes{Name: name},
-	}
-
+func computeMemberId(peerURLs types.URLs, clusterName string, now *time.Time) types.ID {
 	var b []byte
-	sort.Strings(m.PeerURLs)
-	for _, p := range m.PeerURLs {
+	peerURLstrs := peerURLs.StringSlice()
+	sort.Strings(peerURLstrs)
+	for _, p := range peerURLstrs {
 		b = append(b, []byte(p)...)
 	}
 
@@ -79,7 +74,18 @@ func newMember(name string, peerURLs types.URLs, clusterName string, now *time.T
 	}
 
 	hash := sha1.Sum(b)
-	m.ID = types.ID(binary.BigEndian.Uint64(hash[:8]))
+	return types.ID(binary.BigEndian.Uint64(hash[:8]))
+}
+
+func newMember(name string, peerURLs types.URLs, memberId types.ID, isLearner bool) *Member {
+	m := &Member{
+		RaftAttributes: RaftAttributes{
+			PeerURLs:  peerURLs.StringSlice(),
+			IsLearner: isLearner,
+		},
+		Attributes: Attributes{Name: name},
+		ID:         memberId,
+	}
 	return m
 }
 
diff --git a/server/etcdserver/api/membership/store.go b/server/etcdserver/api/membership/store.go
index 95c2e6aaa..c4b5a2d60 100644
--- a/server/etcdserver/api/membership/store.go
+++ b/server/etcdserver/api/membership/store.go
@@ -57,6 +57,14 @@ func mustSaveMemberToBackend(lg *zap.Logger, be backend.Backend, m *Member) {
 	tx.UnsafePut(membersBucketName, mkey, mvalue)
 }
 
+func TrimClusterFromBackend(be backend.Backend) error {
+	tx := be.BatchTx()
+	tx.Lock()
+	defer tx.Unlock()
+	tx.UnsafeDeleteBucket(clusterBucketName)
+	return nil
+}
+
 func mustDeleteMemberFromBackend(be backend.Backend, id types.ID) {
 	mkey := backendMemberKey(id)
 
diff --git a/server/mvcc/backend/batch_tx.go b/server/mvcc/backend/batch_tx.go
index eb75c29fc..d4bc8c684 100644
--- a/server/mvcc/backend/batch_tx.go
+++ b/server/mvcc/backend/batch_tx.go
@@ -28,6 +28,7 @@ import (
 type BatchTx interface {
 	ReadTx
 	UnsafeCreateBucket(name []byte)
+	UnsafeDeleteBucket(name []byte)
 	UnsafePut(bucketName []byte, key []byte, value []byte)
 	UnsafeSeqPut(bucketName []byte, key []byte, value []byte)
 	UnsafeDelete(bucketName []byte, key []byte)
@@ -80,6 +81,18 @@ func (t *batchTx) UnsafeCreateBucket(name []byte) {
 	t.pending++
 }
 
+func (t *batchTx) UnsafeDeleteBucket(name []byte) {
+	err := t.tx.DeleteBucket(name)
+	if err != nil && err != bolt.ErrBucketNotFound {
+		t.backend.lg.Fatal(
+			"failed to delete a bucket",
+			zap.String("bucket-name", string(name)),
+			zap.Error(err),
+		)
+	}
+	t.pending++
+}
+
 // UnsafePut must be called holding the lock on the tx.
 func (t *batchTx) UnsafePut(bucketName []byte, key []byte, value []byte) {
 	t.unsafePut(bucketName, key, value, false)
diff --git a/server/mvcc/kvstore_test.go b/server/mvcc/kvstore_test.go
index 1bb3fae24..10fa50053 100644
--- a/server/mvcc/kvstore_test.go
+++ b/server/mvcc/kvstore_test.go
@@ -875,6 +875,7 @@ func (b *fakeBatchTx) Unlock()                        {}
 func (b *fakeBatchTx) RLock()                         {}
 func (b *fakeBatchTx) RUnlock()                       {}
 func (b *fakeBatchTx) UnsafeCreateBucket(name []byte) {}
+func (b *fakeBatchTx) UnsafeDeleteBucket(name []byte) {}
 func (b *fakeBatchTx) UnsafePut(bucketName []byte, key []byte, value []byte) {
 	b.Recorder.Record(testutil.Action{Name: "put", Params: []interface{}{bucketName, key, value}})
 }