crossks: clean up virtual server info on close and bootstrap failure

Project: pingcap/tidb | Language: Go | Concepts: error-handling, memory-management, distributed-systems, testing-strategy | Difficulty: mid
Link: https://github.com/pingcap/tidb/pull/70519

The problem

When a cross-keyspace runtime in TiDB was closed or failed to bootstrap, its virtual server-info record in etcd was not immediately removed. This meant that stale server-info records could linger for up to 90 seconds (until their lease expired), potentially causing the system to incorrectly assume that the server was still part of the cluster. This could lead to issues like cluster-table fanout including phantom servers, or DDL owners waiting on non-existent nodes.

Why it's non-trivial

A naive fix might just close the etcd client or stop the goroutines, but that would not guarantee immediate removal of the server-info record or revocation of its lease. Additionally, bootstrap failures are tricky because the runtime may not be fully registered with the manager, so normal cleanup paths would not be triggered. Proper cleanup must handle both normal shutdown and partial initialization, and must ensure that etcd records and leases are revoked even if the runtime context has already been canceled.

The approach

1. Track server-info registration and clean up on bootstrap failure

The code now tracks whether the server-info was registered, and if so, ensures it is removed and its lease revoked on error:

+	var svrInfoSyncer *serverinfo.Syncer
+	serverInfoRegistered := false
 	defer func() {
 		if err != nil {
+			if serverInfoRegistered {
+				svrInfoSyncer.RemoveServerInfo()
+				svrInfoSyncer.RevokeSession()
+			}
 			cancel()
 			err2 := etcdCli.Close()
 			if err2 != nil {
 				logger.Warn("failed to close etcd client", zap.Error(err2))
 			}
 		}
 	}()

Annotation:

• serverInfoRegistered is set to true after successful registration.
• If an error occurs after registration, the deferred function removes the server-info record and revokes its lease before canceling the context or closing the etcd client.

2. Register and track the server-info syncer

-	svrInfoSyncer := serverinfo.NewCrossKSSyncer(
+	svrInfoSyncer = serverinfo.NewCrossKSSyncer(
 		virtualSvrID,
 		func() uint64 {
 			// this ID is used to allocate connection ID, since we don't accept
 			// connections, just use 0.
 			return 0
 		},
 		etcdCli,
 		ks,
 	)
 	if err = svrInfoSyncer.NewSessionAndStoreServerInfo(ctx); err != nil {
 		return nil, errors.Trace(err)
 	}
+	serverInfoRegistered = true

Annotation:

• The syncer is now assigned to a variable that outlives this block, so it can be referenced in deferred cleanup.
• serverInfoRegistered is set only after successful registration, so cleanup is only attempted if needed.

3. Remove server-info and revoke lease on normal close

 func (m *SessionManager) close() {
 	close(m.exitCh)
 	m.cancel()
 	m.wg.Wait()
+	if m.svrInfoSyncer != nil {
+		m.svrInfoSyncer.RemoveServerInfo()
+		m.svrInfoSyncer.RevokeSession()
+	}
 	m.schemaVerSyncer.Close()
 	if err := m.etcdCli.Close(); err != nil {
 		logger.Warn("failed to close etcd client", zap.Error(err))
 	}
 }

Annotation:

• On normal shutdown, after all goroutines have exited, the code removes the server-info and revokes the etcd session lease, ensuring no stale records remain.

4. Implement lease revocation in the syncer

// RevokeSession stops refreshing the current server-info session and revokes its lease.
// The caller must stop ServerInfoSyncLoop first so it cannot recreate the session.
func (s *Syncer) RevokeSession() {
	if s.etcdCli == nil || s.session == nil {
		return
	}
	session := s.session
	lease := session.Lease()
	session.Orphan()

	ctx, cancel := context.WithTimeout(context.Background(), KeyOpDefaultTimeout)
	defer cancel()
	if _, err := s.etcdCli.Revoke(ctx, lease); err != nil {
		logutil.BgLogger().Warn("revoke server info lease failed",
			zap.Int64("lease", int64(lease)),
			zap.Stringer("serverInfo", s.GetLocalServerInfo()),
			zap.Error(err))
	}
}

Annotation:

• This function revokes the etcd lease for the server-info session using a bounded context, so it works even after the main context is canceled.
• The session is orphaned first to prevent refresh races.

5. Add regression tests for both cleanup paths

Tests verify that server-info records are immediately removed and leases revoked on both normal close and bootstrap failure:

t.Run("close removes virtual server info", func(t *testing.T) {
	// ...
	userKSDom.GetCrossKSMgr().CloseKS(keyspace.System)
	resp, err = serverInfoObserver.Get(context.Background(), serverInfoKey)
	require.NoError(t, err)
	require.Empty(t, resp.Kvs, "virtual server info should be removed when the runtime closes")
	requireLeaseRevoked(t, serverInfoObserver, serverInfoLease)
})

t.Run("bootstrap failure removes virtual server info", func(t *testing.T) {
	// ...
	require.ErrorContains(t, err, "injected schema syncer init failure")
	// ...
	resp, err := observer.Get(context.Background(), faultKV.serverInfoKey)
	require.NoError(t, err)
	require.Empty(t, resp.Kvs, "failed bootstrap should remove its virtual server info")
	requireLeaseRevoked(t, observer, faultKV.serverInfoLease)
})

Annotation:

• These tests inject failures and check that the etcd record is gone and the lease is revoked immediately, not after the TTL.

From the review thread

No notable review discussion on this PR.

Takeaways

• Always clean up distributed system metadata (like etcd records and leases) immediately on shutdown or failure, not just on process exit or timeout.
• When initialization can partially succeed, track which resources were acquired so you can clean them up in error paths.
• Use bounded contexts for resource cleanup to ensure cleanup actions can complete even if the main context has been canceled.