feat: fix checkMySQL and add tests and documentation (vitessio#11895)

Signed-off-by: Manan Gupta <manan@planetscale.com>

Signed-off-by: Manan Gupta <manan@planetscale.com>

go/vt/vttablet/tabletserver/state_manager.go

@@ -38,10 +38,14 @@ type servingState int64
 
 const (
 	// StateNotConnected is the state where tabletserver is not
-	// connected to an underlying mysql instance.
+	// connected to an underlying mysql instance. In this state we close
+	// query engine since MySQL is probably unavailable
 	StateNotConnected = servingState(iota)
 	// StateNotServing is the state where tabletserver is connected
 	// to an underlying mysql instance, but is not serving queries.
+	// We do not close the query engine to not close the pool. We keep
+	// the query engine open but prevent queries from running by blocking them
+	// in StartRequest.
 	StateNotServing
 	// StateServing is where queries are allowed.
 	StateServing
@@ -330,11 +334,25 @@ func (sm *stateManager) checkMySQL() {
 		}
 		defer sm.transitioning.Release()
 
+		// This is required to prevent new queries from running in StartRequest
+		// unless they are part of a running transaction.
+		sm.setWantState(StateNotConnected)
 		sm.closeAll()
+
+		// Now that we reached the NotConnected state, we want to go back to the
+		// Serving state. The retry will only succeed once MySQL is reachable again
+		// Until then EnsureConnectionAndDB will error out.
+		sm.setWantState(StateServing)
 		sm.retryTransition(fmt.Sprintf("Cannot connect to MySQL, shutting down query service: %v", err))
 	}()
 }
 
+func (sm *stateManager) setWantState(stateWanted servingState) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+	sm.wantState = stateWanted
+}
+
 // isCheckMySQLRunning returns 1 if CheckMySQL function is in progress
 func (sm *stateManager) isCheckMySQLRunning() int64 {
 	if sm.checkMySQLRunning.Get() {

go/vt/vttablet/tabletserver/state_manager_test.go

@@ -456,10 +456,15 @@ func TestStateManagerCheckMySQL(t *testing.T) {
 	err := sm.SetServingType(topodatapb.TabletType_PRIMARY, testNow, StateServing, "")
 	require.NoError(t, err)
 
+	sm.te = &delayedTxEngine{}
 	sm.qe.(*testQueryEngine).failMySQL = true
 	order.Set(0)
 	sm.checkMySQL()
+	// We know checkMySQL will take atleast 50 milliseconds since txEngine.Close has a sleep in the test code
+	time.Sleep(10 * time.Millisecond)
 	assert.EqualValues(t, 1, sm.isCheckMySQLRunning())
+	// When we are in CheckMySQL state, we should not be accepting any new requests which aren't transactional
+	assert.False(t, sm.IsServing())
 
 	// Rechecking immediately should be a no-op:
 	sm.checkMySQL()
@@ -491,6 +496,7 @@ func TestStateManagerCheckMySQL(t *testing.T) {
 		time.Sleep(10 * time.Millisecond)
 	}
 
+	assert.True(t, sm.IsServing())
 	assert.Equal(t, topodatapb.TabletType_PRIMARY, sm.Target().TabletType)
 	assert.Equal(t, StateServing, sm.State())