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

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

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
@@ -325,11 +329,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
+}
+
 // StopService shuts down sm. If the shutdown doesn't complete
 // within timeBombDuration, it crashes the process.
 func (sm *stateManager) StopService() {

go/vt/vttablet/tabletserver/state_manager_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package tabletserver
 
 import (
+ "context"
  "errors"
  "sync"
  "testing"
@@ -26,8 +27,6 @@ import (
 
  "vitess.io/vitess/go/mysql/fakesqldb"
 
- "context"
-
  "github.com/stretchr/testify/assert"
  "github.com/stretchr/testify/require"
 
@@ -457,9 +456,16 @@ 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)
+ // this asserts that checkMySQL is running
+ assert.EqualValues(t, 0, sm.checkMySQLThrottler.Size())
+ // 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 +497,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())
 }