Added ability to keep instance in the ring when BasicLifecycler is shutting down

CHANGELOG.md

@@ -22,5 +22,6 @@
 * [ENHANCEMENT] Optimise memberlist receive path when used as a backing store for rings with a large number of members. #76 #77 #84 #91
 * [ENHANCEMENT] Memberlist: prepare the data to send on the write before starting counting the elapsed time for `-memberlist.packet-write-timeout`, in order to reduce chances we hit the timeout when sending a packet to other node. #89
 * [ENHANCEMENT] flagext: for cases such as `DeprecatedFlag()` that need a logger, add RegisterFlagsWithLogger. #80
+* [ENHANCEMENT] Added option to BasicLifecycler to keep instance in the ring when stopping. #97
 * [BUGFIX] spanlogger: Support multiple tenant IDs. #59
 * [BUGFIX] Memberlist: fixed corrupted packets when sending compound messages with more than 255 messages or messages bigger than 64KB. #85

ring/basic_lifecycler.go

@@ -51,6 +51,10 @@ type BasicLifecyclerConfig struct {
 	HeartbeatPeriod     time.Duration
 	TokensObservePeriod time.Duration
 	NumTokens           int
+
+	// If true lifecycler doesn't unregister instance from the ring when it's stopping. Default value is false,
+	// which means unregistering.
+	KeepInstanceInTheRingOnShutdown bool
 }
 
 // BasicLifecycler is a basic ring lifecycler which allows to hook custom
@@ -227,11 +231,15 @@ heartbeatLoop:
 		}
 	}
 
-	// Remove the instance from the ring.
-	if err := l.unregisterInstance(context.Background()); err != nil {
-		return errors.Wrapf(err, "failed to unregister instance from the ring (ring: %s)", l.ringName)
+	if l.cfg.KeepInstanceInTheRingOnShutdown {
+		level.Info(l.logger).Log("msg", "keeping instance the ring", "ring", l.ringName)
+	} else {
+		// Remove the instance from the ring.
+		if err := l.unregisterInstance(context.Background()); err != nil {
+			return errors.Wrapf(err, "failed to unregister instance from the ring (ring: %s)", l.ringName)
+		}
+		level.Info(l.logger).Log("msg", "instance removed from the ring", "ring", l.ringName)
 	}
-	level.Info(l.logger).Log("msg", "instance removed from the ring", "ring", l.ringName)
 
 	return nil
 }

ring/basic_lifecycler_test.go

@@ -190,6 +190,46 @@ func TestBasicLifecycler_UnregisterOnStop(t *testing.T) {
 	assert.False(t, ok)
 }
 
+func TestBasicLifecycler_KeepInTheRingOnStop(t *testing.T) {
+	ctx := context.Background()
+	cfg := prepareBasicLifecyclerConfig()
+	cfg.KeepInstanceInTheRingOnShutdown = true
+
+	lifecycler, delegate, store, err := prepareBasicLifecycler(t, cfg)
+	require.NoError(t, err)
+
+	delegate.onRegister = func(_ *BasicLifecycler, _ Desc, _ bool, _ string, _ InstanceDesc) (InstanceState, Tokens) {
+		return ACTIVE, Tokens{1, 2, 3, 4, 5}
+	}
+	delegate.onStopping = func(lifecycler *BasicLifecycler) {
+		require.NoError(t, lifecycler.changeState(context.Background(), LEAVING))
+	}
+
+	require.NoError(t, services.StartAndAwaitRunning(ctx, lifecycler))
+	assert.Equal(t, ACTIVE, lifecycler.GetState())
+	assert.Equal(t, Tokens{1, 2, 3, 4, 5}, lifecycler.GetTokens())
+	assert.True(t, lifecycler.IsRegistered())
+	assert.NotZero(t, lifecycler.GetRegisteredAt())
+	assert.Equal(t, float64(cfg.NumTokens), testutil.ToFloat64(lifecycler.metrics.tokensOwned))
+	assert.Equal(t, float64(cfg.NumTokens), testutil.ToFloat64(lifecycler.metrics.tokensToOwn))
+
+	require.NoError(t, services.StopAndAwaitTerminated(ctx, lifecycler))
+	assert.Equal(t, LEAVING, lifecycler.GetState())
+	assert.Equal(t, Tokens{1, 2, 3, 4, 5}, lifecycler.GetTokens())
+	assert.True(t, lifecycler.IsRegistered())
+	assert.NotZero(t, lifecycler.GetRegisteredAt())
+	assert.Equal(t, float64(cfg.NumTokens), testutil.ToFloat64(lifecycler.metrics.tokensOwned))
+	assert.Equal(t, float64(cfg.NumTokens), testutil.ToFloat64(lifecycler.metrics.tokensToOwn))
+
+	// Assert on the instance is in the ring.
+	inst, ok := getInstanceFromStore(t, store, testInstanceID)
+	assert.True(t, ok)
+	assert.Equal(t, cfg.Addr, inst.GetAddr())
+	assert.Equal(t, LEAVING, inst.GetState())
+	assert.Equal(t, Tokens{1, 2, 3, 4, 5}, Tokens(inst.GetTokens()))
+	assert.Equal(t, cfg.Zone, inst.GetZone())
+}
+
 func TestBasicLifecycler_HeartbeatWhileRunning(t *testing.T) {
 	ctx := context.Background()
 	cfg := prepareBasicLifecyclerConfig()