From 780abdff2135d86287dac18a0cd2ea20d1e5e802 Mon Sep 17 00:00:00 2001
From: Arthur Schreiber <arthurschreiber@github.com>
Date: Thu, 30 May 2024 11:48:41 +0000
Subject: [PATCH] more WIP fixes

---
 go/pools/smartconnpool/pool.go                    | 5 ++++-
 go/vt/vttablet/tabletserver/repltracker/poller.go | 6 +++++-
 go/vt/vttablet/tabletserver/repltracker/reader.go | 3 ++-
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/go/pools/smartconnpool/pool.go b/go/pools/smartconnpool/pool.go
index 7c10d6ba4b0..47d80aa3fc8 100644
--- a/go/pools/smartconnpool/pool.go
+++ b/go/pools/smartconnpool/pool.go
@@ -337,7 +337,10 @@ func (pool *ConnPool[C]) put(conn *Pooled[C]) {
 
 	if conn == nil {
 		var err error
-		conn, err = pool.connNew(context.Background())
+		// TODO: Do we really want to wait for up to a second here?
+		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+		defer cancel()
+		conn, err = pool.connNew(ctx)
 		if err != nil {
 			pool.closedConn()
 			return
diff --git a/go/vt/vttablet/tabletserver/repltracker/poller.go b/go/vt/vttablet/tabletserver/repltracker/poller.go
index ace01dffb2d..21a2e9baf1c 100644
--- a/go/vt/vttablet/tabletserver/repltracker/poller.go
+++ b/go/vt/vttablet/tabletserver/repltracker/poller.go
@@ -17,6 +17,7 @@ limitations under the License.
 package repltracker
 
 import (
+	"context"
 	"sync"
 	"time"
 
@@ -45,7 +46,10 @@ func (p *poller) Status() (time.Duration, error) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
-	status, err := p.mysqld.ReplicationStatus()
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	status, err := p.mysqld.ReplicationStatusWithContext(ctx)
 	if err != nil {
 		return 0, err
 	}
diff --git a/go/vt/vttablet/tabletserver/repltracker/reader.go b/go/vt/vttablet/tabletserver/repltracker/reader.go
index fe469bb2e31..985bb136f1f 100644
--- a/go/vt/vttablet/tabletserver/repltracker/reader.go
+++ b/go/vt/vttablet/tabletserver/repltracker/reader.go
@@ -142,7 +142,7 @@ func (r *heartbeatReader) Status() (time.Duration, error) {
 func (r *heartbeatReader) readHeartbeat() {
 	defer r.env.LogError()
 
-	ctx, cancel := context.WithDeadline(context.Background(), r.now().Add(r.interval))
+	ctx, cancel := context.WithTimeout(context.Background(), r.interval)
 	defer cancel()
 
 	res, err := r.fetchMostRecentHeartbeat(ctx)
@@ -150,6 +150,7 @@ func (r *heartbeatReader) readHeartbeat() {
 		r.recordError(vterrors.Wrap(err, "failed to read most recent heartbeat"))
 		return
 	}
+
 	ts, err := parseHeartbeatResult(res)
 	if err != nil {
 		r.recordError(vterrors.Wrap(err, "failed to parse heartbeat result"))