discovery: hold connMu when writing tabletHealthCheck health fields

go/vt/discovery/healthcheck.go

@@ -915,10 +915,14 @@ func (hc *HealthCheckImpl) registeredHealthCheck(alias *topodata.TabletAlias) *t
 // TabletConnection returns the Connection to a given tablet.
 func (hc *HealthCheckImpl) TabletConnection(ctx context.Context, alias *topodata.TabletAlias, target *query.Target) (queryservice.QueryService, error) {
 	thc := hc.registeredHealthCheck(alias)
-	if thc == nil || thc.Conn == nil {
+	if thc == nil {
 		return nil, vterrors.Errorf(vtrpc.Code_NOT_FOUND, "tablet: %v is either down or nonexistent", alias)
 	}
-	return thc.Connection(ctx), nil
+	conn := thc.currentConnection()
+	if conn == nil {
+		return nil, vterrors.Errorf(vtrpc.Code_NOT_FOUND, "tablet: %v is either down or nonexistent (no health check connection)", alias)
+	}
+	return conn, nil
 }
 
 // getAliasByCell should only be called while holding hc.mu

go/vt/discovery/healthcheck_test.go

@@ -37,6 +37,7 @@ import (
 
 	"vitess.io/vitess/go/test/utils"
 	"vitess.io/vitess/go/vt/grpcclient"
+	"vitess.io/vitess/go/vt/logutil"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/memorytopo"
 	"vitess.io/vitess/go/vt/topo/topoproto"
@@ -45,6 +46,7 @@ import (
 	"vitess.io/vitess/go/vt/vttablet/tabletconn"
 	"vitess.io/vitess/go/vt/vttablet/tabletconntest"
 
+	logutilpb "vitess.io/vitess/go/vt/proto/logutil"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 )
@@ -281,6 +283,191 @@ func TestHealthCheck(t *testing.T) {
 	testChecksum(t, 0, hc.stateChecksum())
 }
 
+// TestHealthCheckConcurrentReadDuringUpdate exercises GetTabletHealthByAlias
+// (which copies the tablet's health fields under connMu via SimpleCopy) while
+// the tablet's checkConn goroutine processes streaming health responses. The
+// field writes in processResponse must hold connMu, otherwise this reports a
+// data race under -race.
+func TestHealthCheckConcurrentReadDuringUpdate(t *testing.T) {
+	ctx := utils.LeakCheckContext(t)
+	hcErrorCounters.ResetAll()
+	ts := memorytopo.NewServer(ctx, "cell")
+	defer ts.Close()
+	hc := createTestHc(ctx, ts)
+	defer hc.Close()
+	tablet := createTestTablet(0, "cell", "a")
+	tablet.Type = topodatapb.TabletType_REPLICA
+	input := make(chan *querypb.StreamHealthResponse)
+	_ = createFakeConn(tablet, input)
+	hc.AddTablet(tablet)
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	stop := make(chan struct{})
+
+	// Writer: drive processResponse repeatedly through the checkConn goroutine.
+	go func() {
+		defer wg.Done()
+		for i := 0; ; i++ {
+			shr := &querypb.StreamHealthResponse{
+				TabletAlias:   tablet.Alias,
+				Target:        &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
+				Serving:       i%2 == 0,
+				RealtimeStats: &querypb.RealtimeStats{ReplicationLagSeconds: uint32(i % 5)},
+			}
+			select {
+			case input <- shr:
+			case <-stop:
+				return
+			}
+		}
+	}()
+
+	// Reader: copy the same fields under connMu via SimpleCopy.
+	go func() {
+		defer wg.Done()
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				_, _ = hc.GetTabletHealthByAlias(tablet.Alias)
+			}
+		}
+	}()
+
+	time.Sleep(200 * time.Millisecond)
+	close(stop)
+	wg.Wait()
+}
+
+// TestTabletConnectionConcurrentWithStreamClose exercises TabletConnection,
+// which reads thc.Conn, while the tablet's checkConn goroutine repeatedly
+// closes and re-dials the connection on stream errors. closeConnection and
+// finalizeConn write thc.Conn under connMu, so TabletConnection must read it
+// under connMu too, otherwise this reports a data race under -race.
+func TestTabletConnectionConcurrentWithStreamClose(t *testing.T) {
+	ctx := utils.LeakCheckContext(t)
+	hcErrorCounters.ResetAll()
+	ts := memorytopo.NewServer(ctx, "cell")
+	defer ts.Close()
+	hc := createTestHc(ctx, ts)
+	defer hc.Close()
+	tablet := createTestTablet(0, "cell", "a")
+	tablet.Type = topodatapb.TabletType_REPLICA
+	input := make(chan *querypb.StreamHealthResponse)
+	fc := createFakeConn(tablet, input)
+	fc.errCh = make(chan error)
+	hc.AddTablet(tablet)
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	stop := make(chan struct{})
+
+	// Writer: alternate a health response with a stream error so the checkConn
+	// goroutine repeatedly closes and re-dials the connection.
+	go func() {
+		defer wg.Done()
+		shr := &querypb.StreamHealthResponse{
+			TabletAlias:   tablet.Alias,
+			Target:        &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
+			Serving:       true,
+			RealtimeStats: &querypb.RealtimeStats{ReplicationLagSeconds: 1},
+		}
+		for {
+			select {
+			case input <- shr:
+			case <-stop:
+				return
+			}
+			select {
+			case fc.errCh <- errors.New("some stream error"):
+			case <-stop:
+				return
+			}
+		}
+	}()
+
+	// Reader: read thc.Conn through TabletConnection.
+	go func() {
+		defer wg.Done()
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				_, _ = hc.TabletConnection(ctx, tablet.Alias, nil)
+			}
+		}
+	}()
+
+	time.Sleep(200 * time.Millisecond)
+	close(stop)
+	wg.Wait()
+}
+
+// TestHealthCheckReentrantLoggerCallback ensures the serving-state change log
+// is emitted without holding connMu. The configured logger can be a
+// CallbackLogger whose callback calls back into the healthcheck (e.g. via
+// GetTabletHealthByAlias, whose SimpleCopy takes connMu); if setServingState
+// logged with connMu held, that callback would deadlock the checkConn
+// goroutine and this test would time out.
+func TestHealthCheckReentrantLoggerCallback(t *testing.T) {
+	ctx := utils.LeakCheckContext(t)
+	hcErrorCounters.ResetAll()
+	ts := memorytopo.NewServer(ctx, "cell")
+	defer ts.Close()
+
+	tablet := createTestTablet(0, "cell", "a")
+	tablet.Type = topodatapb.TabletType_REPLICA
+
+	// The healthcheck is created with the logger already set, so the callback
+	// reads it through an atomic pointer stored right after creation.
+	var hcPtr atomic.Pointer[HealthCheckImpl]
+	reentered := make(chan struct{}, 1)
+	logger := logutil.NewCallbackLogger(func(e *logutilpb.Event) {
+		if !strings.Contains(e.Value, "HealthCheckUpdate(Serving State)") {
+			return
+		}
+		hc := hcPtr.Load()
+		if hc == nil {
+			return
+		}
+		_, _ = hc.GetTabletHealthByAlias(tablet.Alias)
+		select {
+		case reentered <- struct{}{}:
+		default:
+		}
+	})
+
+	hc := NewHealthCheck(ctx, 1*time.Millisecond, time.Hour, ts, "cell", "", nil, WithLogger(logger))
+	defer hc.Close()
+	hcPtr.Store(hc)
+
+	input := make(chan *querypb.StreamHealthResponse)
+	_ = createFakeConn(tablet, input)
+	hc.AddTablet(tablet)
+
+	// The first processed response logs the serving-state change from the
+	// checkConn goroutine, which runs the callback above.
+	shr := &querypb.StreamHealthResponse{
+		TabletAlias:   tablet.Alias,
+		Target:        &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
+		Serving:       true,
+		RealtimeStats: &querypb.RealtimeStats{ReplicationLagSeconds: 1},
+	}
+	select {
+	case input <- shr:
+	case <-time.After(30 * time.Second):
+		require.FailNow(t, "timed out sending the health response")
+	}
+	select {
+	case <-reentered:
+	case <-time.After(30 * time.Second):
+		require.FailNow(t, "logger callback did not complete; the serving-state log likely ran while holding connMu")
+	}
+}
+
 func TestHealthCheckStreamError(t *testing.T) {
 	ctx := utils.LeakCheckContext(t)
 

go/vt/discovery/tablet_health_check.go

@@ -48,7 +48,12 @@ type tabletHealthCheck struct {
 	cancelFunc context.CancelFunc
 	// Tablet is the tablet object that was sent to HealthCheck.AddTablet.
 	Tablet *topodata.Tablet
-	// mutex to protect Conn
+	// connMu protects Conn and all the mutable health fields below
+	// (Target, Serving, PrimaryTermStartTime, Stats, LastError): the
+	// checkConn goroutine writes them while readers such as SimpleCopy
+	// and currentConnection copy them from other goroutines. It must
+	// not be held across network IO (connection Close or dialing) or
+	// while invoking the logger.
 	connMu sync.Mutex
 	// Conn is the connection associated with the tablet.
 	Conn queryservice.QueryService
@@ -100,16 +105,19 @@ func (thc *tabletHealthCheck) SimpleCopy() *TabletHealth {
 
 // setServingState sets the tablet state to the given value.
 //
-// If the state changes, it logs the change so that failures
-// from the health check connection are logged the first time,
-// but don't continue to log if the connection stays down.
+// If the state changes, it returns a function that logs the change so that
+// failures from the health check connection are logged the first time, but
+// don't continue to log if the connection stays down. It returns nil when
+// there is nothing to log.
 //
-// thc.mu must be locked before calling this function
-func (thc *tabletHealthCheck) setServingState(serving bool, reason string) {
+// thc.connMu must be locked before calling this function. The returned log
+// function must be called only after connMu has been released: the logger can
+// be a CallbackLogger whose callback re-enters the healthcheck (e.g. via
+// GetTabletHealthByAlias) and would deadlock on connMu.
+func (thc *tabletHealthCheck) setServingState(serving bool, reason string) func() {
+	var logServingChange func()
 	if !thc.loggedServingState || (serving != thc.Serving) {
-		// Emit the log from a separate goroutine to avoid holding
-		// the th lock while logging is happening
-		thc.logger.Infof("HealthCheckUpdate(Serving State): tablet: %v serving %v => %v for %v/%v (%v) reason: %s",
+		msg := fmt.Sprintf("HealthCheckUpdate(Serving State): tablet: %v serving %v => %v for %v/%v (%v) reason: %s",
 			topotools.TabletIdent(thc.Tablet),
 			thc.Serving,
 			serving,
@@ -118,9 +126,13 @@ func (thc *tabletHealthCheck) setServingState(serving bool, reason string) {
 			thc.Target.GetTabletType(),
 			reason,
 		)
+		logServingChange = func() {
+			thc.logger.Infof("%s", msg)
+		}
 		thc.loggedServingState = true
 	}
 	thc.Serving = serving
+	return logServingChange
 }
 
 // stream streams healthcheck responses to callback.
@@ -144,6 +156,16 @@ func (thc *tabletHealthCheck) Connection(ctx context.Context) queryservice.Query
 	return thc.connectionLocked(ctx)
 }
 
+// currentConnection returns the current connection under connMu without
+// attempting to (re)dial. It returns nil when there is no connection. This is
+// used by callers that only need to read thc.Conn, which is written under
+// connMu by closeConnection and finalizeConn.
+func (thc *tabletHealthCheck) currentConnection() queryservice.QueryService {
+	thc.connMu.Lock()
+	defer thc.connMu.Unlock()
+	return thc.Conn
+}
+
 func (thc *tabletHealthCheck) connectionLocked(ctx context.Context) queryservice.QueryService {
 	if thc.Conn == nil {
 		conn, err := tabletconn.GetDialer()(ctx, thc.Tablet, grpcclient.FailFast(true))
@@ -185,6 +207,7 @@ func (thc *tabletHealthCheck) processResponse(hc *HealthCheckImpl, shr *query.St
 		return vterrors.New(vtrpc.Code_FAILED_PRECONDITION, fmt.Sprintf("health stats mismatch, tablet %+v alias does not match response alias %v", thc.Tablet, shr.TabletAlias))
 	}
 
+	thc.connMu.Lock()
 	prevTarget := thc.Target
 	// check whether this is a trivial update so as to update healthy map
 	trivialUpdate := thc.LastError == nil && thc.Serving && shr.RealtimeStats.HealthError == "" && shr.Serving &&
@@ -198,10 +221,14 @@ func (thc *tabletHealthCheck) processResponse(hc *HealthCheckImpl, shr *query.St
 	if healthErr != nil {
 		reason = "healthCheck update error: " + healthErr.Error()
 	}
-	thc.setServingState(serving, reason)
+	logServingChange := thc.setServingState(serving, reason)
+	thc.connMu.Unlock()
+	if logServingChange != nil {
+		logServingChange()
+	}
 
 	// notify downstream for primary change
-	hc.updateHealth(thc, prevTarget, trivialUpdate, thc.Serving)
+	hc.updateHealth(thc, prevTarget, trivialUpdate, serving)
 	return nil
 }
 
@@ -308,12 +335,18 @@ func (thc *tabletHealthCheck) checkConn(hc *HealthCheckImpl) {
 		// This will ensure that this update prevails over any previous message that
 		// stream could have sent.
 		if timedout.Load() {
+			thc.connMu.Lock()
 			thc.LastError = fmt.Errorf("healthcheck timed out (latest %v)", thc.lastResponseTimestamp)
-			thc.setServingState(false, thc.LastError.Error())
-			hcErrorCounters.Add([]string{thc.Target.Keyspace, thc.Target.Shard, topoproto.TabletTypeLString(thc.Target.TabletType)}, 1)
+			logServingChange := thc.setServingState(false, thc.LastError.Error())
+			target := thc.Target
+			thc.connMu.Unlock()
+			if logServingChange != nil {
+				logServingChange()
+			}
+			hcErrorCounters.Add([]string{target.Keyspace, target.Shard, topoproto.TabletTypeLString(target.TabletType)}, 1)
 			// trivialUpdate = false because this is an error
 			// up = false because we did not get a healthy response within the timeout
-			hc.updateHealth(thc, thc.Target, false, false)
+			hc.updateHealth(thc, target, false, false)
 		}
 
 		// Streaming RPC failed e.g. because vttablet was restarted or took too long.
@@ -334,25 +367,37 @@ func (thc *tabletHealthCheck) checkConn(hc *HealthCheckImpl) {
 
 func (thc *tabletHealthCheck) closeConnection(ctx context.Context, err error) {
 	thc.logger.Warningf("tablet %v healthcheck stream error: %v", thc.Tablet, err)
-	thc.setServingState(false, err.Error())
+	thc.connMu.Lock()
+	logServingChange := thc.setServingState(false, err.Error())
 	thc.LastError = err
-	_ = thc.Conn.Close(ctx)
+	conn := thc.Conn
 	thc.Conn = nil
+	thc.connMu.Unlock()
+	if logServingChange != nil {
+		logServingChange()
+	}
+	_ = conn.Close(ctx)
 }
 
 // finalizeConn closes the health checking connection.
 // To be called only on exit from checkConn().
 func (thc *tabletHealthCheck) finalizeConn() {
-	thc.setServingState(false, "finalizeConn closing connection")
+	thc.connMu.Lock()
+	logServingChange := thc.setServingState(false, "finalizeConn closing connection")
 	// Note: checkConn() exits only when thc.ctx.Done() is closed. Thus it's
 	// safe to simply get Err() value here and assign to LastError.
 	thc.LastError = thc.ctx.Err()
-	if thc.Conn != nil {
+	conn := thc.Conn
+	thc.Conn = nil
+	thc.connMu.Unlock()
+	if logServingChange != nil {
+		logServingChange()
+	}
+	if conn != nil {
 		// Don't use thc.ctx because it's already closed.
 		// Use a separate context, and add a timeout to prevent unbounded waits.
 		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 		defer cancel()
-		_ = thc.Conn.Close(ctx)
-		thc.Conn = nil
+		_ = conn.Close(ctx)
 	}
 }