github.com/matrixorigin/matrixone@v0.7.0/pkg/frontend/routine_test.go

// Copyright 2021 Matrix Origin
//
// 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 frontend

import (
	"context"
	"database/sql"
	"fmt"
	"github.com/golang/mock/gomock"
	"github.com/google/uuid"
	"github.com/matrixorigin/matrixone/pkg/config"
	mock_frontend "github.com/matrixorigin/matrixone/pkg/frontend/test"
	"github.com/matrixorigin/matrixone/pkg/logutil"
	"github.com/matrixorigin/matrixone/pkg/sql/parsers"
	"github.com/matrixorigin/matrixone/pkg/sql/parsers/dialect"
	"github.com/matrixorigin/matrixone/pkg/sql/parsers/tree"
	"github.com/matrixorigin/matrixone/pkg/util/metric"
	"github.com/matrixorigin/matrixone/pkg/vm/engine"
	"github.com/matrixorigin/matrixone/pkg/vm/process"
	"github.com/prashantv/gostub"
	pcg "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/sync/errgroup"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func Test_inc_dec(t *testing.T) {
	rt := &Routine{}
	counter := int32(0)
	eg := errgroup.Group{}

	eg.Go(func() error {
		rt.increaseCount(func() {
			atomic.AddInt32(&counter, 1)
		})
		return nil
	})
	time.Sleep(100 * time.Millisecond)

	eg.Go(func() error {
		rt.decreaseCount(func() {
			atomic.AddInt32(&counter, -1)
		})
		return nil
	})
	time.Sleep(100 * time.Millisecond)

	err := eg.Wait()
	assert.NoError(t, err)
	assert.Equal(t, counter, int32(0))
	assert.False(t, rt.connectionBeCounted.Load())
}

type genMrs func(ses *Session) *MysqlResultSet

var newMockWrapper = func(ctrl *gomock.Controller, ses *Session,
	sql2result map[string]genMrs,
	sql2NoResultSet map[string]bool, sql string, stmt tree.Statement, proc *process.Process) ComputationWrapper {
	var mrs *MysqlResultSet
	var columns []interface{}
	var ok, ok2 bool
	var err error
	var gen genMrs
	if gen, ok = sql2result[sql]; ok {
		mrs = gen(ses)
		for _, col := range mrs.Columns {
			columns = append(columns, col)
		}
	} else if _, ok2 = sql2NoResultSet[sql]; ok2 {
		//no result set
	} else {
		panic(fmt.Sprintf("there is no mysqlResultset for the sql %s", sql))
	}
	uuid, _ := uuid.NewUUID()
	runner := mock_frontend.NewMockComputationRunner(ctrl)
	runner.EXPECT().Run(gomock.Any()).DoAndReturn(func(uint64) error {
		proto := ses.GetMysqlProtocol()
		if mrs != nil {
			err = proto.SendResultSetTextBatchRowSpeedup(mrs, mrs.GetRowCount())
			if err != nil {
				logutil.Errorf("flush error %v", err)
				return err
			}
		}
		return nil
	}).AnyTimes()
	mcw := mock_frontend.NewMockComputationWrapper(ctrl)
	mcw.EXPECT().GetAst().Return(stmt).AnyTimes()
	mcw.EXPECT().GetProcess().Return(proc).AnyTimes()
	mcw.EXPECT().SetDatabaseName(gomock.Any()).Return(nil).AnyTimes()
	mcw.EXPECT().GetColumns().Return(columns, nil).AnyTimes()
	mcw.EXPECT().GetAffectedRows().Return(uint64(0)).AnyTimes()
	mcw.EXPECT().Compile(gomock.Any(), gomock.Any(), gomock.Any()).Return(runner, nil).AnyTimes()
	mcw.EXPECT().GetUUID().Return(uuid[:]).AnyTimes()
	mcw.EXPECT().RecordExecPlan(gomock.Any()).Return(nil).AnyTimes()
	mcw.EXPECT().GetLoadTag().Return(false).AnyTimes()
	return mcw
}

func Test_ConnectionCount(t *testing.T) {
	//client connection method: mysql -h 127.0.0.1 -P 6001 --default-auth=mysql_native_password -uroot -p
	//client connect
	//ion method: mysql -h 127.0.0.1 -P 6001 -udump -p
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	var conn1, conn2 *sql.DB
	var err error

	//before anything using the configuration
	eng := mock_frontend.NewMockEngine(ctrl)
	eng.EXPECT().New(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
	eng.EXPECT().Commit(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
	eng.EXPECT().Rollback(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
	txnClient := mock_frontend.NewMockTxnClient(ctrl)
	pu, err := getParameterUnit("test/system_vars_config.toml", eng, txnClient)
	require.NoError(t, err)

	noResultSet := make(map[string]bool)
	resultSet := make(map[string]genMrs)

	var wrapperStubFunc = func(db, sql, user string, eng engine.Engine, proc *process.Process, ses *Session) ([]ComputationWrapper, error) {
		var cw []ComputationWrapper = nil
		var stmts []tree.Statement = nil
		var cmdFieldStmt *InternalCmdFieldList
		var err error
		if isCmdFieldListSql(sql) {
			cmdFieldStmt, err = parseCmdFieldList(proc.Ctx, sql)
			if err != nil {
				return nil, err
			}
			stmts = append(stmts, cmdFieldStmt)
		} else {
			stmts, err = parsers.Parse(proc.Ctx, dialect.MYSQL, sql)
			if err != nil {
				return nil, err
			}
		}

		for _, stmt := range stmts {
			cw = append(cw, newMockWrapper(ctrl, ses, resultSet, noResultSet, sql, stmt, proc))
		}
		return cw, nil
	}

	bhStub := gostub.Stub(&GetComputationWrapper, wrapperStubFunc)
	defer bhStub.Reset()

	ctx := context.WithValue(context.TODO(), config.ParameterUnitKey, pu)
	rm, _ := NewRoutineManager(ctx, pu)
	rm.SetSkipCheckUser(true)

	wg := sync.WaitGroup{}
	wg.Add(1)

	//running server
	go func() {
		defer wg.Done()
		echoServer(rm.Handler, rm, NewSqlCodec())
	}()

	cCounter := metric.ConnectionCounter(sysAccountName)
	cCounter.Set(0)

	time.Sleep(time.Second * 2)
	conn1, err = openDbConn(t, 6001)
	require.NoError(t, err)

	time.Sleep(time.Second * 2)
	conn2, err = openDbConn(t, 6001)
	require.NoError(t, err)

	time.Sleep(time.Second * 2)

	cc := rm.clientCount()
	assert.GreaterOrEqual(t, cc, 2)

	x := &pcg.Metric{}
	err = cCounter.Write(x)
	assert.NoError(t, err)
	assert.GreaterOrEqual(t, x.Gauge.GetValue(), float64(2))

	//close the connection
	closeDbConn(t, conn1)
	closeDbConn(t, conn2)

	time.Sleep(time.Second * 2)

	cc = rm.clientCount()
	assert.GreaterOrEqual(t, cc, 0)

	err = cCounter.Write(x)
	assert.NoError(t, err)
	assert.GreaterOrEqual(t, x.Gauge.GetValue(), float64(0))

	time.Sleep(time.Millisecond * 10)
	//close server
	setServer(1)
	wg.Wait()
}