go/vt/discovery: configurable logger (vitessio#17846)

Signed-off-by: Max Englander <[email protected]>

Author: maxenglander

go/vt/discovery/healthcheck.go

@@ -54,6 +54,7 @@ import (
 	"vitess.io/vitess/go/netutil"
 	"vitess.io/vitess/go/stats"
 	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/logutil"
 	"vitess.io/vitess/go/vt/proto/query"
 	"vitess.io/vitess/go/vt/proto/topodata"
 	"vitess.io/vitess/go/vt/proto/vtrpc"
@@ -190,8 +191,10 @@ func FilteringKeyspaces() bool {
 	return len(KeyspacesToWatch) > 0
 }
 
-type KeyspaceShardTabletType string
-type tabletAliasString string
+type (
+	KeyspaceShardTabletType string
+	tabletAliasString       string
+)
 
 // HealthCheck declares what the TabletGateway needs from the HealthCheck
 type HealthCheck interface {
@@ -301,6 +304,9 @@ type HealthCheckImpl struct {
 	subscribers map[chan *TabletHealth]string
 	// loadTabletsTrigger is used to immediately load information about tablets of a specific shard.
 	loadTabletsTrigger chan topo.KeyspaceShard
+	// options contains optional settings used to modify HealthCheckImpl
+	// behavior.
+	options Options
 }
 
 // NewVTGateHealthCheckFilters returns healthcheck filters for vtgate.
@@ -352,9 +358,9 @@ func NewVTGateHealthCheckFilters() (filters TabletFilters, err error) {
 // filters.
 //
 //	Is one or more filters to apply when determining what tablets we want to stream healthchecks from.
-func NewHealthCheck(ctx context.Context, retryDelay, healthCheckTimeout time.Duration, topoServer *topo.Server, localCell, cellsToWatch string, filters TabletFilter) *HealthCheckImpl {
-	log.Infof("loading tablets for cells: %v", cellsToWatch)
-
+func NewHealthCheck(
+	ctx context.Context, retryDelay, healthCheckTimeout time.Duration, topoServer *topo.Server, localCell, cellsToWatch string, filters TabletFilter, opts ...Option,
+) *HealthCheckImpl {
 	hc := &HealthCheckImpl{
 		ts:                 topoServer,
 		cell:               localCell,
@@ -366,19 +372,23 @@ func NewHealthCheck(ctx context.Context, retryDelay, healthCheckTimeout time.Dur
 		subscribers:        make(map[chan *TabletHealth]string),
 		cellAliases:        make(map[string]string),
 		loadTabletsTrigger: make(chan topo.KeyspaceShard, 1024),
+		options:            withOptions(opts...),
 	}
+
+	hc.logger().Infof("loading tablets for cells: %v", cellsToWatch)
+
 	var topoWatchers []*TopologyWatcher
 	cells := strings.Split(cellsToWatch, ",")
 	if cellsToWatch == "" {
 		cells = append(cells, localCell)
 	}
 
 	for _, c := range cells {
-		log.Infof("Setting up healthcheck for cell: %v", c)
+		hc.logger().Infof("Setting up healthcheck for cell: %v", c)
 		if c == "" {
 			continue
 		}
-		topoWatchers = append(topoWatchers, NewTopologyWatcher(ctx, topoServer, hc, filters, c, refreshInterval, refreshKnownTablets))
+		topoWatchers = append(topoWatchers, NewTopologyWatcher(ctx, topoServer, hc, filters, c, refreshInterval, refreshKnownTablets, opts...))
 	}
 
 	hc.topoWatchers = topoWatchers
@@ -403,7 +413,7 @@ func (hc *HealthCheckImpl) AddTablet(tablet *topodata.Tablet) {
 		return
 	}
 
-	log.Infof("Adding tablet to healthcheck: %v", tablet)
+	hc.logger().Infof("Adding tablet to healthcheck: %v", tablet)
 	hc.mu.Lock()
 	defer hc.mu.Unlock()
 	if hc.healthByAlias == nil {
@@ -421,14 +431,15 @@ func (hc *HealthCheckImpl) AddTablet(tablet *topodata.Tablet) {
 		cancelFunc: cancelFunc,
 		Tablet:     tablet,
 		Target:     target,
+		logger:     hc.logger(),
 	}
 
 	// add to our datastore
 	key := KeyFromTarget(target)
 	tabletAlias := topoproto.TabletAliasString(tablet.Alias)
 	if _, ok := hc.healthByAlias[tabletAliasString(tabletAlias)]; ok {
 		// We should not add a tablet that we already have
-		log.Errorf("Program bug: tried to add existing tablet: %v to healthcheck", tabletAlias)
+		hc.logger().Errorf("Program bug: tried to add existing tablet: %v to healthcheck", tabletAlias)
 		return
 	}
 	hc.healthByAlias[tabletAliasString(tabletAlias)] = thc
@@ -456,7 +467,7 @@ func (hc *HealthCheckImpl) ReplaceTablet(old, new *topodata.Tablet) {
 }
 
 func (hc *HealthCheckImpl) deleteTablet(tablet *topodata.Tablet) {
-	log.Infof("Removing tablet from healthcheck: %v", tablet)
+	hc.logger().Infof("Removing tablet from healthcheck: %v", tablet)
 	hc.mu.Lock()
 	defer hc.mu.Unlock()
 
@@ -499,7 +510,7 @@ func (hc *HealthCheckImpl) deleteTablet(tablet *topodata.Tablet) {
 	// delete from authoritative map
 	th, ok := hc.healthByAlias[tabletAlias]
 	if !ok {
-		log.Infof("We have no health data for tablet: %v, it might have been deleted already", tablet)
+		hc.logger().Infof("We have no health data for tablet: %v, it might have been deleted already", tablet)
 		return
 	}
 	// Calling this will end the context associated with th.checkConn,
@@ -518,7 +529,7 @@ func (hc *HealthCheckImpl) updateHealth(th *TabletHealth, prevTarget *query.Targ
 	// so that we're not racing to update it and in effect re-adding a copy of the
 	// tablet record that was deleted
 	if _, ok := hc.healthByAlias[tabletAlias]; !ok {
-		log.Infof("Tablet %v has been deleted, skipping health update", th.Tablet)
+		hc.logger().Infof("Tablet %v has been deleted, skipping health update", th.Tablet)
 		return
 	}
 
@@ -543,7 +554,7 @@ func (hc *HealthCheckImpl) updateHealth(th *TabletHealth, prevTarget *query.Targ
 		// causing an interruption where no primary is assigned to the shard.
 		if prevTarget.TabletType == topodata.TabletType_PRIMARY {
 			if primaries := hc.healthData[oldTargetKey]; len(primaries) == 0 {
-				log.Infof("We will have no health data for the next new primary tablet after demoting the tablet: %v, so start loading tablets now", topotools.TabletIdent(th.Tablet))
+				hc.logger().Infof("We will have no health data for the next new primary tablet after demoting the tablet: %v, so start loading tablets now", topotools.TabletIdent(th.Tablet))
 				// We want to trigger a call to load tablets for this keyspace-shard,
 				// but we want this to be non-blocking to prevent the code from deadlocking as described in https://github.com/vitessio/vitess/issues/16994.
 				// If the buffer is exhausted, then we'll just receive the update when all the tablets are loaded on the ticker.
@@ -572,7 +583,7 @@ func (hc *HealthCheckImpl) updateHealth(th *TabletHealth, prevTarget *query.Targ
 			// We already have one up server, see if we
 			// need to replace it.
 			if th.PrimaryTermStartTime < hc.healthy[targetKey][0].PrimaryTermStartTime {
-				log.Warningf("not marking healthy primary %s as Up for %s because its PrimaryTermStartTime is smaller than the highest known timestamp from previous PRIMARYs %s: %d < %d ",
+				hc.logger().Warningf("not marking healthy primary %s as Up for %s because its PrimaryTermStartTime is smaller than the highest known timestamp from previous PRIMARYs %s: %d < %d ",
 					topoproto.TabletAliasString(th.Tablet.Alias),
 					topoproto.KeyspaceShardString(th.Target.Keyspace, th.Target.Shard),
 					topoproto.TabletAliasString(hc.healthy[targetKey][0].Tablet.Alias),
@@ -609,7 +620,7 @@ func (hc *HealthCheckImpl) updateHealth(th *TabletHealth, prevTarget *query.Targ
 
 	isNewPrimary := isPrimary && prevTarget.TabletType != topodata.TabletType_PRIMARY
 	if isNewPrimary {
-		log.Errorf("Adding 1 to PrimaryPromoted counter for target: %v, tablet: %v, tabletType: %v", prevTarget, topoproto.TabletAliasString(th.Tablet.Alias), th.Target.TabletType)
+		hc.logger().Errorf("Adding 1 to PrimaryPromoted counter for target: %v, tablet: %v, tabletType: %v", prevTarget, topoproto.TabletAliasString(th.Tablet.Alias), th.Target.TabletType)
 		hcPrimaryPromotedCounters.Add([]string{th.Target.Keyspace, th.Target.Shard}, 1)
 	}
 
@@ -666,7 +677,7 @@ func (hc *HealthCheckImpl) broadcast(th *TabletHealth) {
 		default:
 			// If the channel is full, we drop the message.
 			hcChannelFullCounter.Add(1)
-			log.Warningf("HealthCheck broadcast channel is full for %v, dropping message for %s", subscriber, topotools.TabletIdent(th.Tablet))
+			hc.logger().Warningf("HealthCheck broadcast channel is full for %v, dropping message for %s", subscriber, topotools.TabletIdent(th.Tablet))
 			// Print the stack trace only once.
 			printStack()
 		}
@@ -849,7 +860,7 @@ func (hc *HealthCheckImpl) waitForTablets(ctx context.Context, targets []*query.
 			timer.Stop()
 			for _, target := range targets {
 				if target != nil {
-					log.Infof("couldn't find tablets for target: %v", target)
+					hc.logger().Infof("couldn't find tablets for target: %v", target)
 				}
 			}
 			return ctx.Err()
@@ -977,7 +988,7 @@ func (hc *HealthCheckImpl) ServeHTTP(w http.ResponseWriter, _ *http.Request) {
 	if err != nil {
 		// Error logged
 		if _, err := w.Write([]byte(err.Error())); err != nil {
-			log.Errorf("write to buffer error failed: %v", err)
+			hc.logger().Errorf("write to buffer error failed: %v", err)
 		}
 
 		return
@@ -988,7 +999,7 @@ func (hc *HealthCheckImpl) ServeHTTP(w http.ResponseWriter, _ *http.Request) {
 
 	// Error logged
 	if _, err := w.Write(buf.Bytes()); err != nil {
-		log.Errorf("write to buffer bytes failed: %v", err)
+		hc.logger().Errorf("write to buffer bytes failed: %v", err)
 	}
 }
 
@@ -1029,6 +1040,11 @@ func (hc *HealthCheckImpl) stateChecksum() int64 {
 	return int64(crc32.ChecksumIEEE(buf.Bytes()))
 }
 
+// logger returns the logutil.Logger used by the healthcheck.
+func (hc *HealthCheckImpl) logger() logutil.Logger {
+	return hc.options.logger
+}
+
 // TabletToMapKey creates a key to the map from tablet's host and ports.
 // It should only be used in discovery and related module.
 func TabletToMapKey(tablet *topodata.Tablet) string {

go/vt/discovery/options.go

@@ -0,0 +1,67 @@
+/*
+Copyright 2025 The Vitess Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+	http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package discovery
+
+import (
+	"vitess.io/vitess/go/vt/logutil"
+)
+
+// Options configure a discovery components. Options are set by the Option
+// values passed to the component constructors.
+type Options struct {
+	logger logutil.Logger
+}
+
+// Option configures how we perform certain operations.
+type Option interface {
+	apply(*Options)
+}
+
+// funcOption wraps a function that modifies options into an implementation of
+// the Option interface.
+type funcOption struct {
+	f func(*Options)
+}
+
+func defaultOptions() Options {
+	return Options{
+		logger: logutil.NewConsoleLogger(),
+	}
+}
+
+func withOptions(dos ...Option) Options {
+	os := defaultOptions()
+	for _, do := range dos {
+		do.apply(&os)
+	}
+	return os
+}
+
+func (fhco *funcOption) apply(dos *Options) {
+	fhco.f(dos)
+}
+
+func newFuncOption(f func(*Options)) *funcOption {
+	return &funcOption{
+		f: f,
+	}
+}
+
+// WithLogger accepts a custom logger to use in a discovery component. If this
+// option is not provided then the default system logger will be used.
+func WithLogger(l logutil.Logger) Option {
+	return newFuncOption(func(o *Options) {
+		o.logger = l
+	})
+}

go/vt/discovery/tablet_health_check.go

@@ -25,7 +25,7 @@ import (
 	"time"
 
 	"vitess.io/vitess/go/vt/grpcclient"
-	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/logutil"
 	"vitess.io/vitess/go/vt/proto/vtrpc"
 	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/topotools"
@@ -71,6 +71,8 @@ type tabletHealthCheck struct {
 	// possibly delete both these
 	loggedServingState    bool
 	lastResponseTimestamp time.Time // timestamp of the last healthcheck response
+	// logger is used to log messages.
+	logger logutil.Logger
 }
 
 // String is defined because we want to print a []*tabletHealthCheck array nicely.
@@ -107,7 +109,7 @@ func (thc *tabletHealthCheck) setServingState(serving bool, reason string) {
 	if !thc.loggedServingState || (serving != thc.Serving) {
 		// Emit the log from a separate goroutine to avoid holding
 		// the th lock while logging is happening
-		log.Infof("HealthCheckUpdate(Serving State): tablet: %v serving %v => %v for %v/%v (%v) reason: %s",
+		thc.logger.Infof("HealthCheckUpdate(Serving State): tablet: %v serving %v => %v for %v/%v (%v) reason: %s",
 			topotools.TabletIdent(thc.Tablet),
 			thc.Serving,
 			serving,
@@ -294,7 +296,7 @@ func (thc *tabletHealthCheck) checkConn(hc *HealthCheckImpl) {
 			// the healthcheck cache again via the topology watcher.
 			// WARNING: Under no other circumstances should we be deleting the tablet here.
 			if strings.Contains(err.Error(), "health stats mismatch") {
-				log.Warningf("deleting tablet %v from healthcheck due to health stats mismatch", thc.Tablet)
+				thc.logger.Warningf("deleting tablet %v from healthcheck due to health stats mismatch", thc.Tablet)
 				hc.deleteTablet(thc.Tablet)
 				return
 			}
@@ -331,7 +333,7 @@ func (thc *tabletHealthCheck) checkConn(hc *HealthCheckImpl) {
 }
 
 func (thc *tabletHealthCheck) closeConnection(ctx context.Context, err error) {
-	log.Warningf("tablet %v healthcheck stream error: %v", thc.Tablet, err)
+	thc.logger.Warningf("tablet %v healthcheck stream error: %v", thc.Tablet, err)
 	thc.setServingState(false, err.Error())
 	thc.LastError = err
 	_ = thc.Conn.Close(ctx)