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", alias)
+	}
+	return conn, nil
 }
 
 // getAliasByCell should only be called while holding hc.mu

go/vt/discovery/healthcheck_test.go

@@ -281,6 +281,143 @@ 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)
+	// Unblock the writer if it is parked on a send.
+	select {
+	case <-input:
+	default:
+	}
+	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)
+	// Unblock the writer if it is parked on a send.
+	select {
+	case <-input:
+	default:
+	}
+	select {
+	case <-fc.errCh:
+	default:
+	}
+	wg.Wait()
+}
+
 func TestHealthCheckStreamError(t *testing.T) {
 	ctx := utils.LeakCheckContext(t)
 

go/vt/discovery/tablet_health_check.go

@@ -104,7 +104,7 @@ func (thc *tabletHealthCheck) SimpleCopy() *TabletHealth {
 // from the health check connection are logged the first time,
 // but don't continue to log if the connection stays down.
 //
-// thc.mu must be locked before calling this function
+// thc.connMu must be locked before calling this function.
 func (thc *tabletHealthCheck) setServingState(serving bool, reason string) {
 	if !thc.loggedServingState || (serving != thc.Serving) {
 		// Emit the log from a separate goroutine to avoid holding
@@ -144,6 +144,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 +195,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 &&
@@ -199,9 +210,11 @@ func (thc *tabletHealthCheck) processResponse(hc *HealthCheckImpl, shr *query.St
 		reason = "healthCheck update error: " + healthErr.Error()
 	}
 	thc.setServingState(serving, reason)
+	serving = thc.Serving
+	thc.connMu.Unlock()
 
 	// notify downstream for primary change
-	hc.updateHealth(thc, prevTarget, trivialUpdate, thc.Serving)
+	hc.updateHealth(thc, prevTarget, trivialUpdate, serving)
 	return nil
 }
 
@@ -308,12 +321,15 @@ 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)
+			target := thc.Target
+			thc.connMu.Unlock()
+			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 +350,31 @@ 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.connMu.Lock()
 	thc.setServingState(false, err.Error())
 	thc.LastError = err
-	_ = thc.Conn.Close(ctx)
+	conn := thc.Conn
 	thc.Conn = nil
+	thc.connMu.Unlock()
+	_ = conn.Close(ctx)
 }
 
 // finalizeConn closes the health checking connection.
 // To be called only on exit from checkConn().
 func (thc *tabletHealthCheck) finalizeConn() {
+	thc.connMu.Lock()
 	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 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)
 	}
 }