Add election tests

server/backend/backend.go

@@ -34,7 +34,7 @@ import (
 	memdb "github.com/yorkie-team/yorkie/server/backend/database/memory"
 	"github.com/yorkie-team/yorkie/server/backend/database/mongo"
 	"github.com/yorkie-team/yorkie/server/backend/election"
-	dbelection "github.com/yorkie-team/yorkie/server/backend/election/database"
+	mongoelection "github.com/yorkie-team/yorkie/server/backend/election/mongo"
 	"github.com/yorkie-team/yorkie/server/backend/housekeeping"
 	"github.com/yorkie-team/yorkie/server/backend/sync"
 	memsync "github.com/yorkie-team/yorkie/server/backend/sync/memory"
@@ -107,7 +107,7 @@ func New(
 		return nil, err
 	}
 
-	elector := dbelection.NewElector(hostname, db)
+	elector := mongoelection.NewElector(hostname, db)
 
 	keeping, err := housekeeping.Start(
 		housekeepingConf,

server/backend/database/database.go

@@ -276,4 +276,10 @@ type Database interface {
 		leaseLockName string,
 		leaseDuration gotime.Duration,
 	) error
+
+	// FindLeader returns the leader hostname for the given leaseLockName.
+	FindLeader(
+		ctx context.Context,
+		leaseLockName string,
+	) (*string, error)
 }

server/backend/database/memory/database.go

@@ -1337,6 +1337,11 @@ func (d *DB) RenewLeaderLease(
 	return nil
 }
 
+// FindLeader returns the leader hostname for the given leaseLockName.
+func (d *DB) FindLeader(ctx context.Context, leaseLockName string) (*string, error) {
+	return nil, nil
+}
+
 func newID() types.ID {
 	return types.ID(primitive.NewObjectID().Hex())
 }

server/backend/database/mongo/client.go

@@ -1464,6 +1464,32 @@ func (c *Client) RenewLeaderLease(
 	return nil
 }
 
+// FindLeader returns the leader hostname for the given leaseLockName.
+func (c *Client) FindLeader(ctx context.Context, leaseLockName string) (*string, error) {
+	electionInfo := &struct {
+		ElectionId    string      `bson:"election_id"`
+		LeaderID      string      `bson:"leader_id"`
+		LeaseExpireAt gotime.Time `bson:"lease_expire_at"`
+	}{}
+
+	result := c.collection(colElections).FindOne(ctx, bson.M{
+		"election_id": leaseLockName,
+	})
+	if result.Err() == mongo.ErrNoDocuments {
+		return nil, nil
+	}
+	if result.Err() != nil {
+		logging.From(ctx).Error(result.Err())
+		return nil, fmt.Errorf("find leader: %w", result.Err())
+	}
+
+	if err := result.Decode(&electionInfo); err != nil {
+		return nil, fmt.Errorf("decode leader: %w", err)
+	}
+
+	return &electionInfo.LeaderID, nil
+}
+
 func (c *Client) collection(
 	name string,
 	opts ...*options.CollectionOptions,

server/backend/election/database/election.go → server/backend/election/mongo/election.go

@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-// Package database is a database implementation of election package.
-package database
+// Package mongo is a mongo based implementation of election package.
+package mongo
 
 import (
 	"context"
@@ -84,7 +84,7 @@ func (e *Elector) run(
 	onStoppedLeading func(),
 ) {
 	for {
-		ctx := context.Background()
+		ctx, cancelFunc := context.WithCancel(e.ctx)
 		acquired, err := e.database.TryToAcquireLeaderLease(ctx, e.hostname, leaseLockName, leaseDuration)
 		if err != nil {
 			continue
@@ -105,6 +105,7 @@ func (e *Elector) run(
 				select {
 				case <-time.After(leaseDuration / 2):
 				case <-e.ctx.Done():
+					cancelFunc()
 					return
 				}
 			}
@@ -118,6 +119,7 @@ func (e *Elector) run(
 		select {
 		case <-time.After(leaseDuration):
 		case <-e.ctx.Done():
+			cancelFunc()
 			return
 		}
 	}

server/backend/election/mongo/election_test.go

@@ -0,0 +1,98 @@
+package mongo_test
+
+import (
+	"context"
+	"github.com/stretchr/testify/assert"
+	"github.com/yorkie-team/yorkie/server/backend/database/mongo"
+	mongoelection "github.com/yorkie-team/yorkie/server/backend/election/mongo"
+	"github.com/yorkie-team/yorkie/test/helper"
+	"testing"
+	"time"
+)
+
+var (
+	normalTask = func(ctx context.Context) {}
+	stopTask   = func() {}
+)
+
+func setupTestWithDummyData(t *testing.T) *mongo.Client {
+	config := &mongo.Config{
+		ConnectionTimeout: "5s",
+		ConnectionURI:     "mongodb://localhost:27017",
+		YorkieDatabase:    helper.TestDBName(),
+		PingTimeout:       "5s",
+	}
+	assert.NoError(t, config.Validate())
+
+	db, err := mongo.Dial(config)
+	assert.NoError(t, err)
+
+	return db
+}
+
+func TestElection(t *testing.T) {
+	db := setupTestWithDummyData(t)
+
+	t.Run("leader election with multiple electors test", func(t *testing.T) {
+		leaseLockName := t.Name()
+
+		electorA := mongoelection.NewElector("A", db)
+		electorB := mongoelection.NewElector("B", db)
+		electorC := mongoelection.NewElector("C", db)
+
+		assert.NoError(t, electorA.StartElection(leaseLockName, helper.LeaseDuration, normalTask, stopTask))
+		assert.NoError(t, electorB.StartElection(leaseLockName, helper.LeaseDuration, normalTask, stopTask))
+		assert.NoError(t, electorC.StartElection(leaseLockName, helper.LeaseDuration, normalTask, stopTask))
+
+		// elector A will be the leader because it is the first to start the election.
+		leader, err := db.FindLeader(context.Background(), leaseLockName)
+		assert.NoError(t, err)
+
+		assert.Equal(t, "A", *leader)
+
+		// wait for lease expiration and check the leader again
+		// elector A is still the leader because it has renewed the lease.
+		time.Sleep(helper.LeaseDuration)
+		assert.Equal(t, "A", *leader)
+
+		// stop electorA and electorB, then wait for the next leader election
+		// elector C will be the leader because other electors are stopped.
+		assert.NoError(t, electorA.Stop())
+		assert.NoError(t, electorB.Stop())
+
+		time.Sleep(helper.LeaseDuration)
+
+		leader, err = db.FindLeader(context.Background(), leaseLockName)
+		assert.NoError(t, err)
+
+		assert.Equal(t, "C", *leader)
+	})
+
+	t.Run("lease renewal while handling a a long task test", func(t *testing.T) {
+		leaseLockName := t.Name()
+		longTask := func(ctx context.Context) {
+			time.Sleep(helper.LeaseDuration * 2)
+		}
+
+		electorA := mongoelection.NewElector("A", db)
+		electorB := mongoelection.NewElector("B", db)
+		electorC := mongoelection.NewElector("C", db)
+
+		assert.NoError(t, electorA.StartElection(leaseLockName, helper.LeaseDuration, longTask, stopTask))
+		assert.NoError(t, electorB.StartElection(leaseLockName, helper.LeaseDuration, normalTask, stopTask))
+		assert.NoError(t, electorC.StartElection(leaseLockName, helper.LeaseDuration, normalTask, stopTask))
+
+		// check if elector A is still the leader
+		time.Sleep(helper.LeaseDuration)
+
+		leader, err := db.FindLeader(context.Background(), leaseLockName)
+		assert.NoError(t, err)
+
+		assert.Equal(t, "A", *leader)
+	})
+
+	t.Run("handle background routines when shutting down the server test", func(t *testing.T) {
+		// TODO(krapie): find the way to gradually close election routines
+		t.Skip()
+	})
+}

test/helper/helper.go

@@ -74,6 +74,8 @@ var (
 	MongoConnectionURI     = "mongodb://localhost:27017"
 	MongoConnectionTimeout = "5s"
 	MongoPingTimeout       = "5s"
+
+	LeaseDuration = 2 * gotime.Second
 )
 
 func init() {