vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletserver/stateful_connection_pool_test.go (about)

     1  /*
     2  Copyright 2020 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package tabletserver
    18  
    19  import (
    20  	"context"
    21  	"testing"
    22  
    23  	"github.com/stretchr/testify/assert"
    24  	"github.com/stretchr/testify/require"
    25  
    26  	"vitess.io/vitess/go/mysql/fakesqldb"
    27  	"vitess.io/vitess/go/sqltypes"
    28  	querypb "vitess.io/vitess/go/vt/proto/query"
    29  	"vitess.io/vitess/go/vt/vttablet/tabletserver/tx"
    30  )
    31  
    32  var ctx = context.Background()
    33  
    34  func TestActivePoolClientRowsFound(t *testing.T) {
    35  	db := fakesqldb.New(t)
    36  	defer db.Close()
    37  	db.AddQuery("begin", &sqltypes.Result{})
    38  
    39  	pool := newActivePool()
    40  	pool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
    41  
    42  	startNormalSize := pool.conns.Available()
    43  	startFoundRowsSize := pool.foundRowsPool.Available()
    44  
    45  	conn1, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
    46  	require.NoError(t, err)
    47  	assert.Equal(t, startNormalSize-1, pool.conns.Available(), "default pool not used")
    48  
    49  	conn2, err := pool.NewConn(ctx, &querypb.ExecuteOptions{ClientFoundRows: true}, nil)
    50  	require.NoError(t, err)
    51  	assert.Equal(t, startFoundRowsSize-1, pool.conns.Available(), "foundRows pool not used")
    52  
    53  	conn1.Release(tx.TxClose)
    54  	assert.Equal(t, startNormalSize, pool.conns.Available(), "default pool not restored after release")
    55  
    56  	conn2.Release(tx.TxClose)
    57  	assert.Equal(t, startFoundRowsSize, pool.conns.Available(), "default pool not restored after release")
    58  }
    59  
    60  func TestActivePoolForAllTxProps(t *testing.T) {
    61  	db := fakesqldb.New(t)
    62  	defer db.Close()
    63  	pool := newActivePool()
    64  	pool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
    65  	conn1, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
    66  	require.NoError(t, err)
    67  	conn1.txProps = &tx.Properties{}
    68  
    69  	conn2, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
    70  	require.NoError(t, err)
    71  	// for the second connection, we are not going to set a tx state
    72  
    73  	conn3, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
    74  	require.NoError(t, err)
    75  	conn3.txProps = &tx.Properties{}
    76  
    77  	pool.ForAllTxProperties(func(p *tx.Properties) {
    78  		p.LogToFile = true
    79  	})
    80  
    81  	require.True(t, conn1.txProps.LogToFile, "connection missed")
    82  	require.Nil(t, conn2.txProps)
    83  	require.True(t, conn3.txProps.LogToFile, "connection missed")
    84  }
    85  
    86  func TestStatefulPoolShutdownNonTx(t *testing.T) {
    87  	db := fakesqldb.New(t)
    88  	defer db.Close()
    89  	pool := newActivePool()
    90  	pool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
    91  
    92  	// conn1 non-tx, not in use.
    93  	conn1, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
    94  	require.NoError(t, err)
    95  	conn1.Taint(ctx, nil)
    96  	conn1.Unlock()
    97  
    98  	// conn2 tx, not in use.
    99  	conn2, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
   100  	require.NoError(t, err)
   101  	conn2.Taint(ctx, nil)
   102  	conn2.txProps = &tx.Properties{}
   103  	conn2.Unlock()
   104  
   105  	// conn3 non-tx, in use.
   106  	conn3, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
   107  	require.NoError(t, err)
   108  	conn3.Taint(ctx, nil)
   109  
   110  	// After ShutdownNonTx, conn1 should be closed, but not conn3.
   111  	pool.ShutdownNonTx()
   112  	assert.Equal(t, int64(2), pool.active.Size())
   113  	assert.True(t, conn1.IsClosed())
   114  	assert.False(t, conn3.IsClosed())
   115  
   116  	// conn3 should get closed on Unlock.
   117  	conn3.Unlock()
   118  	assert.True(t, conn3.IsClosed())
   119  
   120  	// conn2 should be unaffected.
   121  	assert.False(t, conn2.IsClosed())
   122  }
   123  
   124  func TestStatefulPoolShutdownAll(t *testing.T) {
   125  	db := fakesqldb.New(t)
   126  	defer db.Close()
   127  	pool := newActivePool()
   128  	pool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
   129  
   130  	// conn1 not in use
   131  	conn1, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
   132  	require.NoError(t, err)
   133  	conn1.txProps = &tx.Properties{}
   134  	conn1.Unlock()
   135  
   136  	// conn2 in use.
   137  	conn2, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
   138  	require.NoError(t, err)
   139  	conn2.txProps = &tx.Properties{}
   140  
   141  	conns := pool.ShutdownAll()
   142  	wantconns := []*StatefulConnection{conn1}
   143  	assert.Equal(t, wantconns, conns)
   144  
   145  	// conn2 should get closed on Unlock.
   146  	conn2.Unlock()
   147  	assert.True(t, conn2.IsClosed())
   148  }
   149  
   150  func TestActivePoolGetConnNonExistentTransaction(t *testing.T) {
   151  	db := fakesqldb.New(t)
   152  	defer db.Close()
   153  	pool := newActivePool()
   154  	pool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
   155  	_, err := pool.GetAndLock(12345, "for query")
   156  	require.EqualError(t, err, "not found")
   157  }
   158  
   159  func TestExecWithAbortedCtx(t *testing.T) {
   160  	ctx, cancel := context.WithCancel(ctx)
   161  	db := fakesqldb.New(t)
   162  	defer db.Close()
   163  	pool := newActivePool()
   164  	pool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
   165  	conn, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
   166  	require.NoError(t, err)
   167  	cancel()
   168  	_, err = conn.Exec(ctx, "", 0, false)
   169  	require.Error(t, err)
   170  }
   171  
   172  func TestExecWithDbconnClosed(t *testing.T) {
   173  	db := fakesqldb.New(t)
   174  	defer db.Close()
   175  	pool := newActivePool()
   176  	pool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
   177  	conn, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
   178  	require.NoError(t, err)
   179  	conn.Close()
   180  
   181  	_, err = conn.Exec(ctx, "", 0, false)
   182  	require.EqualError(t, err, "connection was aborted")
   183  }
   184  
   185  func TestExecWithDbconnClosedHavingTx(t *testing.T) {
   186  	db := fakesqldb.New(t)
   187  	defer db.Close()
   188  	pool := newActivePool()
   189  	pool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
   190  	conn, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
   191  	require.NoError(t, err)
   192  	conn.txProps = &tx.Properties{Conclusion: "foobar"}
   193  	conn.Close()
   194  
   195  	_, err = conn.Exec(ctx, "", 0, false)
   196  	require.EqualError(t, err, "transaction was aborted: foobar")
   197  }
   198  
   199  func TestFailOnConnectionRegistering(t *testing.T) {
   200  	db := fakesqldb.New(t)
   201  	defer db.Close()
   202  	pool := newActivePool()
   203  	pool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
   204  	conn, err := pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
   205  	require.NoError(t, err)
   206  	defer conn.Close()
   207  
   208  	pool.lastID.Set(conn.ConnID - 1)
   209  
   210  	_, err = pool.NewConn(ctx, &querypb.ExecuteOptions{}, nil)
   211  	require.Error(t, err, "already present")
   212  }
   213  
   214  func newActivePool() *StatefulConnectionPool {
   215  	env := newEnv("ActivePoolTest")
   216  
   217  	return NewStatefulConnPool(env)
   218  }