From 50f311ac108ecd38c1314ac8153f905a1fefd509 Mon Sep 17 00:00:00 2001
From: Manan Gupta <manan@planetscale.com>
Date: Thu, 19 Dec 2024 18:30:48 +0530
Subject: [PATCH] test: add test for verifying we mark demote primary stalled
 if it takes too longs

Signed-off-by: Manan Gupta <manan@planetscale.com>
---
 .../tabletmanager/rpc_replication_test.go     | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/go/vt/vttablet/tabletmanager/rpc_replication_test.go b/go/vt/vttablet/tabletmanager/rpc_replication_test.go
index c587f1e24b8..a8f2a114f5e 100644
--- a/go/vt/vttablet/tabletmanager/rpc_replication_test.go
+++ b/go/vt/vttablet/tabletmanager/rpc_replication_test.go
@@ -22,6 +22,10 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/semaphore"
+
+	"vitess.io/vitess/go/vt/topo"
+	"vitess.io/vitess/go/vt/vttablet/tabletserver"
 )
 
 // TestWaitForGrantsToHaveApplied tests that waitForGrantsToHaveApplied only succeeds after waitForDBAGrants has been called.
@@ -42,3 +46,50 @@ func TestWaitForGrantsToHaveApplied(t *testing.T) {
 	err = tm.waitForGrantsToHaveApplied(secondContext)
 	require.NoError(t, err)
 }
+
+type demotePrimaryStallQS struct {
+	tabletserver.Controller
+	waitTime       time.Duration
+	primaryStalled bool
+}
+
+func (d *demotePrimaryStallQS) SetDemotePrimaryStalled() {
+	d.primaryStalled = true
+}
+
+func (d *demotePrimaryStallQS) IsServing() bool {
+	time.Sleep(d.waitTime)
+	return false
+}
+
+// TestDemotePrimaryStalled checks that if demote primary takes too long, then we mark it as stalled.
+func TestDemotePrimaryStalled(t *testing.T) {
+	// Set remote operation timeout to a very low value.
+	origVal := topo.RemoteOperationTimeout
+	topo.RemoteOperationTimeout = 100 * time.Millisecond
+	defer func() {
+		topo.RemoteOperationTimeout = origVal
+	}()
+
+	// Create a fake query service control to intercept calls from DemotePrimary function.
+	qsc := &demotePrimaryStallQS{
+		waitTime:       2 * time.Second,
+		primaryStalled: false,
+	}
+	// Create a tablet manager with a replica type tablet.
+	tm := &TabletManager{
+		actionSema:  semaphore.NewWeighted(1),
+		MysqlDaemon: newTestMysqlDaemon(t, 1),
+		tmState: &tmState{
+			displayState: displayState{
+				tablet: newTestTablet(t, 100, "ks", "-", map[string]string{}),
+			},
+		},
+		QueryServiceControl: qsc,
+	}
+
+	// We make IsServing stall for over 2 seconds, which is longer than 10 * remote operation timeout.
+	// This should cause the demote primary operation to be stalled.
+	tm.demotePrimary(context.Background(), false)
+	require.True(t, qsc.primaryStalled)
+}