github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/dbs/session_pool.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package dbs
    15  
    16  import (
    17  	"sync"
    18  
    19  	"github.com/ngaut/pools"
    20  	"github.com/whtcorpsinc/errors"
    21  	"github.com/whtcorpsinc/BerolinaSQL/allegrosql"
    22  	"github.com/whtcorpsinc/milevadb/stochastikctx"
    23  	"github.com/whtcorpsinc/milevadb/soliton/logutil"
    24  	"github.com/whtcorpsinc/milevadb/soliton/mock"
    25  )
    26  
    27  // stochastikPool is used to new stochastik.
    28  type stochastikPool struct {
    29  	mu struct {
    30  		sync.Mutex
    31  		closed bool
    32  	}
    33  	resPool *pools.ResourcePool
    34  }
    35  
    36  func newStochastikPool(resPool *pools.ResourcePool) *stochastikPool {
    37  	return &stochastikPool{resPool: resPool}
    38  }
    39  
    40  // get gets stochastikctx from context resource pool.
    41  // Please remember to call put after you finished using stochastikctx.
    42  func (sg *stochastikPool) get() (stochastikctx.Context, error) {
    43  	if sg.resPool == nil {
    44  		return mock.NewContext(), nil
    45  	}
    46  
    47  	sg.mu.Lock()
    48  	if sg.mu.closed {
    49  		sg.mu.Unlock()
    50  		return nil, errors.Errorf("stochastikPool is closed.")
    51  	}
    52  	sg.mu.Unlock()
    53  
    54  	// no need to protect sg.resPool
    55  	resource, err := sg.resPool.Get()
    56  	if err != nil {
    57  		return nil, errors.Trace(err)
    58  	}
    59  
    60  	ctx := resource.(stochastikctx.Context)
    61  	ctx.GetStochastikVars().SetStatusFlag(allegrosql.ServerStatusAutocommit, true)
    62  	ctx.GetStochastikVars().InRestrictedALLEGROSQL = true
    63  	return ctx, nil
    64  }
    65  
    66  // put returns stochastikctx to context resource pool.
    67  func (sg *stochastikPool) put(ctx stochastikctx.Context) {
    68  	if sg.resPool == nil {
    69  		return
    70  	}
    71  
    72  	// no need to protect sg.resPool, even the sg.resPool is closed, the ctx still need to
    73  	// put into resPool, because when resPool is closing, it will wait all the ctx returns, then resPool finish closing.
    74  	sg.resPool.Put(ctx.(pools.Resource))
    75  }
    76  
    77  // close clean up the stochastikPool.
    78  func (sg *stochastikPool) close() {
    79  	sg.mu.Lock()
    80  	defer sg.mu.Unlock()
    81  	// prevent closing resPool twice.
    82  	if sg.mu.closed || sg.resPool == nil {
    83  		return
    84  	}
    85  	logutil.BgLogger().Info("[dbs] closing stochastikPool")
    86  	sg.resPool.Close()
    87  	sg.mu.closed = true
    88  }