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 }