From 23e9f1d996fcf8821bc1afe26d3670ed911af2d5 Mon Sep 17 00:00:00 2001 From: Kevin Park Date: Sat, 8 Jul 2023 21:10:49 +0900 Subject: [PATCH] Add election test for background routine handling --- server/backend/database/mongo/client.go | 2 +- server/backend/database/mongo/client_test.go | 2 + server/backend/election/mongo/election.go | 10 ++++- .../backend/election/mongo/election_test.go | 42 +++++++++++++++---- 4 files changed, 46 insertions(+), 10 deletions(-) diff --git a/server/backend/database/mongo/client.go b/server/backend/database/mongo/client.go index 1d22af781..609769e0e 100644 --- a/server/backend/database/mongo/client.go +++ b/server/backend/database/mongo/client.go @@ -1467,7 +1467,7 @@ func (c *Client) RenewLeaderLease( // 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"` + ElectionID string `bson:"election_id"` LeaderID string `bson:"leader_id"` LeaseExpireAt gotime.Time `bson:"lease_expire_at"` }{} diff --git a/server/backend/database/mongo/client_test.go b/server/backend/database/mongo/client_test.go index ed26def73..bc6c0f4c7 100644 --- a/server/backend/database/mongo/client_test.go +++ b/server/backend/database/mongo/client_test.go @@ -17,7 +17,9 @@ package mongo_test import ( + "context" "testing" + "time" "github.com/stretchr/testify/assert" diff --git a/server/backend/election/mongo/election.go b/server/backend/election/mongo/election.go index a2858cdd6..237e6ebb6 100644 --- a/server/backend/election/mongo/election.go +++ b/server/backend/election/mongo/election.go @@ -19,6 +19,7 @@ package mongo import ( "context" + "sync" "time" "github.com/yorkie-team/yorkie/server/backend/database" @@ -33,6 +34,8 @@ type Elector struct { ctx context.Context cancelFunc context.CancelFunc + + wg sync.WaitGroup } // NewElector creates a new elector instance. @@ -70,6 +73,7 @@ func (e *Elector) StartElection( // Stop stops all leader elections. func (e *Elector) Stop() error { e.cancelFunc() + e.wg.Wait() return nil } @@ -91,7 +95,11 @@ func (e *Elector) run( } if acquired { - go onStartLeading(ctx) + go func() { + e.wg.Add(1) + onStartLeading(ctx) + e.wg.Done() + }() logging.From(ctx).Infof( "leader elected: %s", e.hostname, ) diff --git a/server/backend/election/mongo/election_test.go b/server/backend/election/mongo/election_test.go index 240b535ab..02c3693d0 100644 --- a/server/backend/election/mongo/election_test.go +++ b/server/backend/election/mongo/election_test.go @@ -2,12 +2,14 @@ package mongo_test import ( "context" + "testing" + "time" + "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 ( @@ -41,8 +43,11 @@ func TestElection(t *testing.T) { electorC := mongoelection.NewElector("C", db) assert.NoError(t, electorA.StartElection(leaseLockName, helper.LeaseDuration, normalTask, stopTask)) + time.Sleep(helper.LeaseDuration) + assert.NoError(t, electorB.StartElection(leaseLockName, helper.LeaseDuration, normalTask, stopTask)) assert.NoError(t, electorC.StartElection(leaseLockName, helper.LeaseDuration, normalTask, stopTask)) + time.Sleep(helper.LeaseDuration) // elector A will be the leader because it is the first to start the election. leader, err := db.FindLeader(context.Background(), leaseLockName) @@ -71,7 +76,7 @@ func TestElection(t *testing.T) { 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) + time.Sleep(helper.LeaseDuration * 4) } electorA := mongoelection.NewElector("A", db) @@ -79,10 +84,12 @@ func TestElection(t *testing.T) { 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)) + time.Sleep(helper.LeaseDuration) + + assert.NoError(t, electorB.StartElection(leaseLockName, helper.LeaseDuration, longTask, stopTask)) + assert.NoError(t, electorC.StartElection(leaseLockName, helper.LeaseDuration, longTask, stopTask)) - // check if elector A is still the leader + // wait for lease expiration and check if elector A is still the leader while handling a long task time.Sleep(helper.LeaseDuration) leader, err := db.FindLeader(context.Background(), leaseLockName) @@ -92,7 +99,26 @@ func TestElection(t *testing.T) { }) 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() + shutdownCh := make(chan struct{}) + + isTaskDone := false + longTask := func(ctx context.Context) { + close(shutdownCh) + time.Sleep(helper.LeaseDuration) + isTaskDone = true + } + + elector := mongoelection.NewElector("A", db) + assert.NoError(t, elector.StartElection(t.Name(), helper.LeaseDuration, longTask, stopTask)) + + // if receive shutdown signal, stop elector + select { + case <-shutdownCh: + assert.NoError(t, elector.Stop()) + } + + // check if the task is done + // this means that the background routine is handled properly after server(elector) is stopped + assert.Equal(t, true, isTaskDone) }) }