Graceful stop database heartbeats in ICS

Support stopping an individual database heartbeat in the inventory
control stream instead of waiting for multiple keepalives to fail.

Author: GavinFrazar

api/client/inventory.go

@@ -343,6 +343,10 @@ func (i *downstreamICS) runSendLoop(stream proto.AuthService_InventoryControlStr
 				oneOf.Msg = &proto.UpstreamInventoryOneOf_Goodbye{
 					Goodbye: msg,
 				}
+			case *proto.UpstreamInventoryStopHeartbeat:
+				oneOf.Msg = &proto.UpstreamInventoryOneOf_StopHeartbeat{
+					StopHeartbeat: msg,
+				}
 			default:
 				sendMsg.errC <- trace.BadParameter("cannot send unexpected upstream msg type: %T", msg)
 				continue
@@ -484,6 +488,8 @@ func (i *upstreamICS) runRecvLoop(stream proto.AuthService_InventoryControlStrea
 			msg = oneOf.GetAgentMetadata()
 		case oneOf.GetGoodbye() != nil:
 			msg = oneOf.GetGoodbye()
+		case oneOf.GetStopHeartbeat() != nil:
+			msg = oneOf.GetStopHeartbeat()
 		default:
 			slog.WarnContext(stream.Context(), "received unknown upstream message", "message", oneOf)
 			continue

api/client/proto/inventory.pb.go

@@ -107,6 +107,8 @@ func (a *UpstreamInventoryAgentMetadata) sealedUpstreamInventoryMessage() {}
 
 func (h *UpstreamInventoryGoodbye) sealedUpstreamInventoryMessage() {}
 
+func (h *UpstreamInventoryStopHeartbeat) sealedUpstreamInventoryMessage() {}
+
 // DownstreamInventoryMessage is a sealed interface representing the possible
 // downstream messages of the inventory controls stream after initial hello.
 type DownstreamInventoryMessage interface {

api/client/proto/types.go

@@ -38,6 +38,9 @@ message UpstreamInventoryOneOf {
     UpstreamInventoryAgentMetadata AgentMetadata = 4;
     // UpstreamInventoryGoodbye advertises that the instance is terminating.
     UpstreamInventoryGoodbye Goodbye = 5;
+    // UpstreamInventoryStopHeartbeat informs the upstream service that the
+    // heartbeat is stopping.
+    UpstreamInventoryStopHeartbeat StopHeartbeat = 6;
   }
 }
 
@@ -166,6 +169,8 @@ message DownstreamInventoryHello {
     bool KubernetesHeartbeats = 17;
     // KubernetesCleanup indicates the ICS supports deleting kubernetes clusters when UpstreamInventoryGoodbye.DeleteResources is set.
     bool KubernetesCleanup = 18;
+    // DatabaseHeartbeatGracefulStop indicates the ICS supports stopping an individual database heartbeat.
+    bool DatabaseHeartbeatGracefulStop = 19;
   }
 
   // SupportedCapabilities advertises the supported features of the auth server.
@@ -255,3 +260,22 @@ message InventoryStatusSummary {
   // ServiceCounts aggregates the number of services.
   map<string, uint32> ServiceCounts = 5;
 }
+
+// UpstreamInventoryStopHeartbeat informs the upstream service that the
+// heartbeat is stopping.
+message UpstreamInventoryStopHeartbeat {
+  // Kind is the kind of heartbeat to stop.
+  StopHeartbeatKind Kind = 1;
+  // HostID is the heartbeat host ID.
+  string HostID = 2;
+  // Name is the name of the heatbeat to stop.
+  string Name = 3;
+}
+
+// StopHeartbeatKind is the type of heartbeat to stop.
+enum StopHeartbeatKind {
+  STOP_HEARTBEAT_KIND_UNSPECIFIED = 0;
+
+  // STOP_HEARTBEAT_KIND_DATABASE_SERVER means stop a database server heartbeat.
+  STOP_HEARTBEAT_KIND_DATABASE_SERVER = 1;
+}

api/proto/teleport/legacy/client/proto/inventory.proto

@@ -4958,13 +4958,14 @@ func (a *Server) RegisterInventoryControlStream(ics client.UpstreamInventoryCont
 		Version:  teleport.Version,
 		ServerID: a.ServerID,
 		Capabilities: &proto.DownstreamInventoryHello_SupportedCapabilities{
-			NodeHeartbeats:       true,
-			AppHeartbeats:        true,
-			AppCleanup:           true,
-			DatabaseHeartbeats:   true,
-			DatabaseCleanup:      true,
-			KubernetesHeartbeats: true,
-			KubernetesCleanup:    true,
+			NodeHeartbeats:                true,
+			AppHeartbeats:                 true,
+			AppCleanup:                    true,
+			DatabaseHeartbeats:            true,
+			DatabaseHeartbeatGracefulStop: true,
+			DatabaseCleanup:               true,
+			KubernetesHeartbeats:          true,
+			KubernetesCleanup:             true,
 		},
 	}
 	if err := ics.Send(a.CloseContext(), downstreamHello); err != nil {

lib/auth/auth.go

@@ -94,6 +94,9 @@ const (
 	dbUpsertRetryOk  testEvent = "db-upsert-retry-ok"
 	dbUpsertRetryErr testEvent = "db-upsert-retry-err"
 
+	dbDelOk  testEvent = "db-del-ok"
+	dbDelErr testEvent = "db-del-err"
+
 	kubeKeepAliveOk  testEvent = "kube-keep-alive-ok"
 	kubeKeepAliveErr testEvent = "kube-keep-alive-err"
 	kubeKeepAliveDel testEvent = "kube-keep-alive-del"
@@ -554,6 +557,21 @@ func (c *Controller) handleControlStream(handle *upstreamHandle) {
 				c.handlePong(handle, m)
 			case *proto.UpstreamInventoryGoodbye:
 				handle.setGoodbye(m)
+			case *proto.UpstreamInventoryStopHeartbeat:
+				key := resourceKey{hostID: m.HostID, name: m.Name}
+				switch m.Kind {
+				case proto.StopHeartbeatKind_STOP_HEARTBEAT_KIND_DATABASE_SERVER:
+					if err := c.handleStopDatabaseServerHB(handle, key); err != nil {
+						handle.CloseWithError(err)
+						return
+					}
+				default:
+					slog.WarnContext(c.closeContext, "Unexpected upstream stop heartbeat kind on control stream",
+						"server_id", handle.Hello().ServerID,
+						"kind", logutils.StringerAttr(m.Kind),
+					)
+				}
+
 			default:
 				slog.WarnContext(c.closeContext, "Unexpected upstream message type on control stream",
 					"message_type", logutils.TypeAttr(m),
@@ -1294,6 +1312,43 @@ func (c *Controller) keepAliveSSHServer(handle *upstreamHandle, now time.Time) e
 	return nil
 }
 
+func (c *Controller) handleStopDatabaseServerHB(handle *upstreamHandle, key resourceKey) error {
+	// the auth layer verifies that a stream's hello message matches the identity and capabilities of the
+	// client cert. after that point it is our responsibility to ensure that heartbeated information is
+	// consistent with the identity and capabilities claimed in the initial hello.
+	if !handle.HasService(types.RoleDatabase) {
+		return trace.AccessDenied("control stream not configured to support database server heartbeats")
+	}
+	if key.hostID != handle.Hello().ServerID {
+		return trace.AccessDenied("incorrect database server ID (expected %q, got %q)", handle.Hello().ServerID, key.hostID)
+	}
+
+	if _, ok := handle.databaseServers[key]; ok {
+		c.testEvent(dbKeepAliveDel)
+		c.onDisconnectFunc(constants.KeepAliveDatabase, 1)
+		if c.dbHBVariableDuration != nil {
+			c.dbHBVariableDuration.Dec()
+		}
+		delete(handle.databaseServers, key)
+		handle.dbKeepAliveDelay.Remove(key)
+	}
+
+	err := c.auth.DeleteDatabaseServer(c.closeContext, apidefaults.Namespace, key.hostID, key.name)
+	if err == nil {
+		c.testEvent(dbDelOk)
+	} else {
+		c.testEvent(dbDelErr)
+		if !trace.IsNotFound(err) {
+			slog.WarnContext(c.closeContext, "Failed to delete database server heartbeat",
+				"database_name", key.name,
+				"server_id", handle.Hello().ServerID,
+				"error", err,
+			)
+		}
+	}
+	return nil
+}
+
 func (c *Controller) createKeepAliveDelay(variableDuration *interval.VariableDuration) *delay.Delay {
 	return delay.New(delay.Params{
 		FirstInterval:    retryutils.HalfJitter(c.serverKeepAlive),

lib/inventory/controller.go

@@ -55,9 +55,11 @@ type fakeAuth struct {
 	mu             sync.Mutex
 	failUpserts    int
 	failKeepAlives int
+	failDeletes    int
 
 	upserts    int
 	keepalives int
+	deletes    int
 	err        error
 
 	expectAddr      string
@@ -125,6 +127,14 @@ func (a *fakeAuth) UpsertDatabaseServer(_ context.Context, server types.Database
 }
 
 func (a *fakeAuth) DeleteDatabaseServer(ctx context.Context, namespace, hostID, name string) error {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	a.deletes++
+
+	if a.failDeletes > 0 {
+		a.failDeletes--
+		return trace.Errorf("delete failed as test condition")
+	}
 	return nil
 }
 
@@ -688,7 +698,7 @@ func TestDatabaseServerBasics(t *testing.T) {
 	require.Equal(t, int64(1), controller.instanceHBVariableDuration.Count())
 
 	// send a fake db server heartbeat
-	for i := 0; i < dbCount; i++ {
+	for i := range dbCount {
 		err := downstream.Send(ctx, &proto.InventoryHeartbeat{
 			DatabaseServer: &types.DatabaseServerV3{
 				Metadata: types.Metadata{
@@ -717,6 +727,53 @@ func TestDatabaseServerBasics(t *testing.T) {
 		deny(dbUpsertErr, dbKeepAliveErr, handlerClose),
 	)
 
+	// set up to induce delete failure
+	auth.mu.Lock()
+	auth.failDeletes = 2
+	auth.mu.Unlock()
+
+	// stop a heartbeat
+	err := downstream.Send(ctx, &proto.UpstreamInventoryStopHeartbeat{
+		Kind:   proto.StopHeartbeatKind_STOP_HEARTBEAT_KIND_DATABASE_SERVER,
+		HostID: serverID,
+		Name:   "db-1",
+	})
+	require.NoError(t, err)
+
+	// verify that keep alive stops, even if the heartbeat couldn't be deleted
+	awaitEvents(t, events,
+		expect(dbKeepAliveDel, dbDelErr),
+		deny(dbDelOk, dbUpsertErr, dbKeepAliveErr, handlerClose),
+	)
+	require.Equal(t, dbCount-1, rc.count())
+
+	// verify heartbeat stop keepalive idempotency
+	err = downstream.Send(ctx, &proto.UpstreamInventoryStopHeartbeat{
+		Kind:   proto.StopHeartbeatKind_STOP_HEARTBEAT_KIND_DATABASE_SERVER,
+		HostID: serverID,
+		Name:   "db-1",
+	})
+	require.NoError(t, err)
+	require.Equal(t, dbCount-1, rc.count())
+
+	awaitEvents(t, events,
+		expect(dbDelErr),
+		deny(dbKeepAliveDel, dbDelOk, dbUpsertErr, dbKeepAliveErr, handlerClose),
+	)
+	require.Equal(t, dbCount-1, rc.count())
+
+	err = downstream.Send(ctx, &proto.UpstreamInventoryStopHeartbeat{
+		Kind:   proto.StopHeartbeatKind_STOP_HEARTBEAT_KIND_DATABASE_SERVER,
+		HostID: serverID,
+		Name:   "db-1",
+	})
+	require.NoError(t, err)
+	awaitEvents(t, events,
+		expect(dbDelOk),
+		deny(dbKeepAliveDel, dbDelErr, dbUpsertErr, dbKeepAliveErr, handlerClose),
+	)
+	require.Equal(t, dbCount-1, rc.count())
+
 	// set up to induce some failures, but not enough to cause the control
 	// stream to be closed.
 	auth.mu.Lock()
@@ -778,7 +835,7 @@ func TestDatabaseServerBasics(t *testing.T) {
 	defer cancel()
 
 	// execute ping
-	_, err := handle.Ping(pingCtx, 1)
+	_, err = handle.Ping(pingCtx, 1)
 	require.NoError(t, err)
 
 	// ensure that local db keepalive states have reset to healthy by waiting
@@ -1741,6 +1798,7 @@ func deny(events ...testEvent) eventOption {
 }
 
 func awaitEvents(t *testing.T, ch <-chan testEvent, opts ...eventOption) {
+	t.Helper()
 	options := eventOpts{
 		expect: make(map[testEvent]int),
 		deny:   make(map[testEvent]struct{}),

lib/inventory/controller_test.go

@@ -35,6 +35,7 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/api/client/proto"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
 	"github.com/gravitational/teleport/api/types"
 	apievents "github.com/gravitational/teleport/api/types/events"
@@ -673,8 +674,8 @@ func (s *Server) stopDynamicLabels(name string) {
 func (s *Server) registerDatabase(ctx context.Context, database types.Database) error {
 	if err := s.startDatabase(ctx, database); err != nil {
 		// Cleanup in case database was initialized only partially.
-		if errStop := s.stopDatabase(ctx, database); errStop != nil {
-			return trace.NewAggregate(err, errStop)
+		if errUnregister := s.unregisterDatabase(ctx, database); errUnregister != nil {
+			return trace.NewAggregate(err, errUnregister)
 		}
 		return trace.Wrap(err)
 	}
@@ -691,10 +692,6 @@ func (s *Server) updateDatabase(ctx context.Context, database types.Database) er
 		return trace.Wrap(err)
 	}
 	if err := s.registerDatabase(ctx, database); err != nil {
-		// If we failed to re-register, don't keep proxying the old database.
-		if errUnregister := s.unregisterDatabase(ctx, database); errUnregister != nil {
-			return trace.NewAggregate(err, errUnregister)
-		}
 		return trace.Wrap(err)
 	}
 	return nil
@@ -721,10 +718,14 @@ func (s *Server) stopProxyingAndDeleteDatabase(ctx context.Context, database typ
 	if err := s.stopDatabase(ctx, database); err != nil {
 		return trace.Wrap(err)
 	}
-	// Heartbeat is stopped but if we don't remove this database server,
-	// it can linger for up to ~10m until its TTL expires.
-	if err := s.deleteDatabaseServer(ctx, database.GetName()); err != nil {
-		return trace.Wrap(err)
+	// stopping the upstream inventory heartbeat may incur a backend write
+	// to delete the heartbeat, which is unnecessary when updating a DB,
+	// since we will upsert a new heartbeat
+	if err := s.stopUpstreamInventoryHeartbeat(ctx, database); err != nil {
+		s.log.WarnContext(ctx, "Failed to stop upstream inventory database heartbeat",
+			"db", log.StringerAttr(database),
+			"error", err,
+		)
 	}
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -781,9 +782,12 @@ func (s *Server) copyDatabaseWithUpdatedLabelsLocked(database types.Database) *t
 func (s *Server) startHeartbeat(ctx context.Context, database types.Database) error {
 	heartbeat, err := srv.NewDatabaseServerHeartbeat(srv.HeartbeatV2Config[*types.DatabaseServerV3]{
 		InventoryHandle: s.cfg.InventoryHandle,
-		Announcer:       s.cfg.AccessPoint,
-		GetResource:     s.getServerInfoFunc(database),
-		OnHeartbeat:     s.cfg.OnHeartbeat,
+		// Announcer is provided to allow falling back to non-ICS heartbeats if
+		// the Auth server is older than the db service.
+		// TODO(gavin): DELETE IN 19.0.0
+		Announcer:   s.cfg.AccessPoint,
+		GetResource: s.getServerInfoFunc(database),
+		OnHeartbeat: s.cfg.OnHeartbeat,
 	})
 	if err != nil {
 		return trace.Wrap(err)
@@ -807,6 +811,37 @@ func (s *Server) stopHeartbeat(name string) error {
 	return heartbeat.Close()
 }
 
+// stopUpstreamInventoryHeartbeat stops the upstream inventory control stream
+// heartbeat for a single database.
+// https://github.com/gravitational/teleport/issues/50237
+func (s *Server) stopUpstreamInventoryHeartbeat(ctx context.Context, db types.Database) error {
+	if _, ok := s.cfg.InventoryHandle.GetSender(); ok {
+		select {
+		// get latest sender
+		case sender := <-s.cfg.InventoryHandle.Sender():
+			if sender.Hello().Capabilities.DatabaseHeartbeatGracefulStop {
+				err := sender.Send(ctx, &proto.UpstreamInventoryStopHeartbeat{
+					Kind:   proto.StopHeartbeatKind_STOP_HEARTBEAT_KIND_DATABASE_SERVER,
+					HostID: s.cfg.HostID,
+					Name:   db.GetName(),
+				})
+				return trace.Wrap(err)
+			}
+		case <-ctx.Done():
+			return trace.Wrap(ctx.Err())
+		}
+	}
+	// the heartbeat may have been created using a fallback method or the
+	// upstream doesn't support graceful stop.
+	// if we don't remove this database server, it can linger for up to ~10m
+	// until its TTL expires.
+	// TODO(gavin): DELETE IN 20.0.0
+	if err := s.deleteDatabaseServer(ctx, db.GetName()); err != nil {
+		return trace.Wrap(err)
+	}
+	return nil
+}
+
 // getServerInfoFunc returns function that the heartbeater uses to report the
 // provided database to the auth server.
 //