119 lines
3.3 KiB
Go
119 lines
3.3 KiB
Go
// Copyright 2023 PingCAP, Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package session_test
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/ngaut/pools"
|
|
"github.com/pingcap/tidb/pkg/ddl/session"
|
|
"github.com/pingcap/tidb/pkg/testkit"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestSessionPool(t *testing.T) {
|
|
store := testkit.CreateMockStore(t)
|
|
tk := testkit.NewTestKit(t, store)
|
|
tk.MustExec("use test")
|
|
resourcePool := pools.NewResourcePool(func() (pools.Resource, error) {
|
|
newTk := testkit.NewTestKit(t, store)
|
|
return newTk.Session(), nil
|
|
}, 4, 4, 0)
|
|
pool := session.NewSessionPool(resourcePool)
|
|
sessCtx, err := pool.Get()
|
|
require.NoError(t, err)
|
|
se := session.NewSession(sessCtx)
|
|
err = se.Begin(context.Background())
|
|
startTS := se.GetSessionVars().TxnCtx.StartTS
|
|
require.NoError(t, err)
|
|
rows, err := se.Execute(context.Background(), "select 2;", "test")
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, len(rows))
|
|
require.Equal(t, int64(2), rows[0].GetInt64(0))
|
|
mgr := tk.Session().GetSessionManager()
|
|
tsList := mgr.GetInternalSessionStartTSList()
|
|
var targetTS uint64
|
|
for _, ts := range tsList {
|
|
if ts == startTS {
|
|
targetTS = ts
|
|
break
|
|
}
|
|
}
|
|
require.NotEqual(t, uint64(0), targetTS)
|
|
err = se.Commit(context.Background())
|
|
pool.Put(sessCtx)
|
|
require.NoError(t, err)
|
|
tsList = mgr.GetInternalSessionStartTSList()
|
|
targetTS = 0
|
|
for _, ts := range tsList {
|
|
if ts == startTS {
|
|
targetTS = ts
|
|
break
|
|
}
|
|
}
|
|
require.Equal(t, uint64(0), targetTS)
|
|
}
|
|
|
|
func TestPessimisticTxn(t *testing.T) {
|
|
store := testkit.CreateMockStore(t)
|
|
tk := testkit.NewTestKit(t, store)
|
|
tk.MustExec("use test")
|
|
tk.MustExec("create table t (a int primary key, b int)")
|
|
tk.MustExec("insert into t values (1, 1)")
|
|
|
|
resourcePool := pools.NewResourcePool(func() (pools.Resource, error) {
|
|
newTk := testkit.NewTestKit(t, store)
|
|
return newTk.Session(), nil
|
|
}, 4, 4, 0)
|
|
pool := session.NewSessionPool(resourcePool)
|
|
ctx := context.Background()
|
|
|
|
sessCtx, err := pool.Get()
|
|
require.NoError(t, err)
|
|
se := session.NewSession(sessCtx)
|
|
sessCtx2, err := pool.Get()
|
|
require.NoError(t, err)
|
|
se2 := session.NewSession(sessCtx2)
|
|
|
|
err = se.BeginPessimistic(ctx)
|
|
require.NoError(t, err)
|
|
err = se2.BeginPessimistic(ctx)
|
|
require.NoError(t, err)
|
|
_, err = se.Execute(ctx, "update test.t set b = b + 1 where a = 1", "ut")
|
|
require.NoError(t, err)
|
|
done := make(chan struct{}, 1)
|
|
go func() {
|
|
_, err := se2.Execute(ctx, "update test.t set b = b + 1 where a = 1", "ut")
|
|
require.NoError(t, err)
|
|
done <- struct{}{}
|
|
err = se2.Commit(ctx)
|
|
require.NoError(t, err)
|
|
close(done)
|
|
}()
|
|
|
|
time.Sleep(100 * time.Millisecond)
|
|
// because this is a pessimistic transaction, the second transaction should be blocked
|
|
require.Len(t, done, 0)
|
|
err = se.Commit(ctx)
|
|
require.NoError(t, err)
|
|
<-done
|
|
_, ok := <-done
|
|
require.False(t, ok)
|
|
pool.Put(sessCtx)
|
|
pool.Put(sessCtx2)
|
|
}
|