Correct health streamer behavior

Signed-off-by: Matt Lord <mattalord@gmail.com>
vitessio · Jul 22, 2023 · ca220b0 · ca220b0
1 parent d5499e4
commit ca220b0
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 11 deletions.
diff --git a/go/vt/discovery/tablet_picker.go b/go/vt/discovery/tablet_picker.go
@@ -18,6 +18,7 @@ package discovery
 
 import (
 	"fmt"
+	"io"
 	"math/rand"
 	"sort"
 	"strings"
@@ -431,13 +432,13 @@ func (tp *TabletPicker) GetMatchingTablets(ctx context.Context) []*topo.TabletIn
 				// Ensure that the tablet is healthy and serving.
 				if err := conn.StreamHealth(ctx, func(shr *querypb.StreamHealthResponse) error {
 					if shr.RealtimeStats.HealthError == "" && shr.Serving {
-						return nil
+						return io.EOF // End the stream
 					}
 					return vterrors.New(vtrpcpb.Code_INTERNAL, "tablet is not serving")
-				}); err == nil {
-					_ = conn.Close(ctx)
+				}); err == nil || err == io.EOF {
 					tablets = append(tablets, tabletInfo)
 				}
+				_ = conn.Close(ctx)
 			}
 		}
 	}

diff --git a/go/vt/vttablet/tabletmanager/vreplication/controller.go b/go/vt/vttablet/tabletmanager/vreplication/controller.go
@@ -19,6 +19,7 @@ package vreplication
 import (
 	"context"
 	"fmt"
+	"io"
 	"strconv"
 	"strings"
 	"sync/atomic"
@@ -336,13 +337,13 @@ func (ct *controller) sourceTabletIsUnhealthy() bool {
 		// Ensure that the tablet is healthy and serving.
 		if err := conn.StreamHealth(ctx, func(shr *querypb.StreamHealthResponse) error {
 			if shr.RealtimeStats.HealthError == "" && shr.Serving {
-				return nil
+				return io.EOF // End the stream
 			}
 			return vterrors.New(vtrpcpb.Code_INTERNAL, "tablet is not serving")
-		}); err == nil {
-			_ = conn.Close(ctx)
+		}); err == nil || err == io.EOF {
 			return true
 		}
+		_ = conn.Close(ctx)
 	}
 	return false
 }

diff --git a/go/vt/vttablet/tabletserver/health_streamer.go b/go/vt/vttablet/tabletserver/health_streamer.go
@@ -181,7 +181,6 @@ func (hs *healthStreamer) Stream(ctx context.Context, callback func(*querypb.Str
 				}
 				return err
 			}
-			return nil
 		}
 	}
 }

diff --git a/go/vt/wrangler/fake_tablet_test.go b/go/vt/wrangler/fake_tablet_test.go
@@ -255,7 +255,6 @@ func (ft *fakeTablet) StreamHealth(ctx context.Context, callback func(*querypb.S
 		Target: &querypb.Target{
 			Keyspace:   ft.Tablet.Keyspace,
 			Shard:      ft.Tablet.Shard,
-			Cell:       ft.Tablet.Alias.Cell,
 			TabletType: ft.Tablet.Type,
 		},
 		RealtimeStats: &querypb.RealtimeStats{},

diff --git a/go/vt/wrangler/traffic_switcher_env_test.go b/go/vt/wrangler/traffic_switcher_env_test.go
@@ -155,7 +155,7 @@ func newTestTableMigraterCustom(ctx context.Context, t *testing.T, sourceShards,
 		}
 		return nil, nil
 	})
-	tabletconntest.SetProtocol("go.vt.traffic_switcher_env_test", dialerName)
+	tabletconntest.SetProtocol("go.vt.wrangler.traffic_switcher_env_test", dialerName)
 
 	vs := &vschemapb.Keyspace{
 		Sharded: true,
@@ -336,7 +336,7 @@ func newTestTablePartialMigrater(ctx context.Context, t *testing.T, shards, shar
 		}
 		return nil, nil
 	})
-	tabletconntest.SetProtocol("go.vt.traffic_switcher_env_test", dialerName)
+	tabletconntest.SetProtocol("go.vt.wrangler.traffic_switcher_env_test", dialerName)
 
 	vs := &vschemapb.Keyspace{
 		Sharded: true,
@@ -494,7 +494,7 @@ func newTestShardMigrater(ctx context.Context, t *testing.T, sourceShards, targe
 		}
 		return nil, nil
 	})
-	tabletconntest.SetProtocol("go.vt.traffic_switcher_env_test", dialerName)
+	tabletconntest.SetProtocol("go.vt.wrangler.traffic_switcher_env_test", dialerName)
 
 	vs := &vschemapb.Keyspace{
 		Sharded: true,