/*

Copyright 2017 Google LLC

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 spanner

import (

"container/heap"

"container/list"

"context"

"fmt"

"log"

"math"

"math/rand"

"runtime/debug"

"strings"

"sync"

"time"

"cloud.google.com/go/internal/trace"

sppb "cloud.google.com/go/spanner/apiv1/spannerpb"

"cloud.google.com/go/spanner/internal"

"go.opencensus.io/stats"

"go.opencensus.io/tag"

octrace "go.opencensus.io/trace"

"go.opentelemetry.io/otel/attribute"

"go.opentelemetry.io/otel/metric"

"google.golang.org/grpc/codes"

"google.golang.org/grpc/metadata"

)

const (

healthCheckIntervalMins = 50

multiplexSessionRefreshInterval = 7 * 24 * time.Hour

)

// ActionOnInactiveTransactionKind describes the kind of action taken when there are inactive transactions.

type ActionOnInactiveTransactionKind int

const (

actionUnspecified ActionOnInactiveTransactionKind = iota

// NoAction action does not perform any action on inactive transactions.

NoAction

// Warn action logs inactive transactions. Any inactive transaction gets logged only once.

Warn

// Close action closes inactive transactions without logging.

Close

// WarnAndClose action logs and closes the inactive transactions.

WarnAndClose

)

// InactiveTransactionRemovalOptions has configurations for action on long-running transactions.

type InactiveTransactionRemovalOptions struct {

// ActionOnInactiveTransaction is the configuration to choose action for inactive transactions.

// It can be one of Warn, Close, WarnAndClose.

ActionOnInactiveTransaction ActionOnInactiveTransactionKind

// long-running transactions will be cleaned up if utilisation is

// greater than the below value.

usedSessionsRatioThreshold float64

// A transaction is considered to be idle if it has not been used for

// a duration greater than the below value.

idleTimeThreshold time.Duration

// frequency for closing inactive transactions

executionFrequency time.Duration

// variable that keeps track of the last execution time when inactive transactions

// were removed by the maintainer task.

lastExecutionTime time.Time

}

// sessionHandle is an interface for transactions to access Cloud Spanner

// sessions safely. It is generated by sessionPool.take().

type sessionHandle struct {

// mu guarantees that the inner session object is returned / destroyed only

// once.

mu sync.Mutex

// session is a pointer to a session object. Transactions never need to

// access it directly.

session *session

// client is the RPC channel to Cloud Spanner. It is set only once during session acquisition.

client spannerClient

// checkoutTime is the time the session was checked out of the pool.

checkoutTime time.Time

// lastUseTime is the time the session was last used after checked out of the pool.

lastUseTime time.Time

// trackedSessionHandle is the linked list node which links the session to

// the list of tracked session handles. trackedSessionHandle is only set if

// TrackSessionHandles has been enabled in the session pool configuration.

trackedSessionHandle *list.Element

// stack is the call stack of the goroutine that checked out the session

// from the pool. This can be used to track down session leak problems.

stack []byte

// eligibleForLongRunning tells if the inner session is eligible to be long-running.

eligibleForLongRunning bool

// if the inner session object is long-running then the stack gets logged once.

isSessionLeakLogged bool

}

// recycle gives the inner session object back to its home session pool. It is

// safe to call recycle multiple times but only the first one would take effect.

func (sh *sessionHandle) recycle() {

sh.mu.Lock()

if sh.session == nil {

// sessionHandle has already been recycled.

sh.mu.Unlock()

return

}

p := sh.session.pool

tracked := sh.trackedSessionHandle

s := sh.session

sh.session = nil

sh.client = nil

sh.trackedSessionHandle = nil

sh.checkoutTime = time.Time{}

sh.lastUseTime = time.Time{}

sh.stack = nil

sh.mu.Unlock()

s.recycle()

if tracked != nil {

p.mu.Lock()

p.trackedSessionHandles.Remove(tracked)

p.mu.Unlock()

}

}

// getID gets the Cloud Spanner session ID from the internal session object.

// getID returns empty string if the sessionHandle is nil or the inner session

// object has been released by recycle / destroy.

func (sh *sessionHandle) getID() string {

sh.mu.Lock()

defer sh.mu.Unlock()

if sh.session == nil {

// sessionHandle has already been recycled/destroyed.

return ""

}

return sh.session.getID()

}

// getClient gets the Cloud Spanner RPC client associated with the session ID

// in sessionHandle.

func (sh *sessionHandle) getClient() spannerClient {

sh.mu.Lock()

defer sh.mu.Unlock()

if sh.session == nil {

return nil

}

if sh.client != nil {

// Use the gRPC connection from the session handle

return sh.client

}

return sh.session.client

}

// getMetadata returns the metadata associated with the session in sessionHandle.

func (sh *sessionHandle) getMetadata() metadata.MD {

sh.mu.Lock()

defer sh.mu.Unlock()

if sh.session == nil {

return nil

}

return sh.session.md

}

// getTransactionID returns the transaction id in the session if available.

func (sh *sessionHandle) getTransactionID() transactionID {

sh.mu.Lock()

defer sh.mu.Unlock()

if sh.session == nil {

return nil

}

return sh.session.tx

}

// destroy destroys the inner session object. It is safe to call destroy

// multiple times and only the first call would attempt to

// destroy the inner session object.

func (sh *sessionHandle) destroy() {

sh.mu.Lock()

s := sh.session

if s == nil {

// sessionHandle has already been recycled.

sh.mu.Unlock()

return

}

tracked := sh.trackedSessionHandle

sh.session = nil

sh.client = nil

sh.trackedSessionHandle = nil

sh.checkoutTime = time.Time{}

sh.lastUseTime = time.Time{}

sh.stack = nil

sh.mu.Unlock()

if tracked != nil {

p := s.pool

p.mu.Lock()

p.trackedSessionHandles.Remove(tracked)

p.mu.Unlock()

}

// since sessionHandle is always used by Transactions we can safely destroy the session with wasInUse=true

s.destroy(false, true)

}

func (sh *sessionHandle) updateLastUseTime() {

sh.mu.Lock()

defer sh.mu.Unlock()

if sh.session != nil {

sh.lastUseTime = time.Now()

}

}

// session wraps a Cloud Spanner session ID through which transactions are

// created and executed.

type session struct {

// client is the RPC channel to Cloud Spanner. It is set only once during

// session's creation.

client spannerClient

// id is the unique id of the session in Cloud Spanner. It is set only once

// during session's creation.

id string

// pool is the session's home session pool where it was created. It is set

// only once during session's creation.

pool *sessionPool

// createTime is the timestamp of the session's creation. It is set only

// once during session's creation.

createTime time.Time

// logger is the logger configured for the Spanner client that created the

// session. If nil, logging will be directed to the standard logger.

logger *log.Logger

// mu protects the following fields from concurrent access: both

// healthcheck workers and transactions can modify them.

mu sync.Mutex

// valid marks the validity of a session.

valid bool

// hcIndex is the index of the session inside the global healthcheck queue.

// If hcIndex < 0, session has been unregistered from the queue.

hcIndex int

// idleList is the linkedlist node which links the session to its home

// session pool's idle list. If idleList == nil, the

// session is not in idle list.

idleList *list.Element

// nextCheck is the timestamp of next scheduled healthcheck of the session.

// It is maintained by the global health checker.

nextCheck time.Time

// checkingHelath is true if currently this session is being processed by

// health checker. Must be modified under health checker lock.

checkingHealth bool

// md is the Metadata to be sent with each request.

md metadata.MD

// tx contains the transaction id if the session has been prepared for

// write.

tx transactionID

// firstHCDone indicates whether the first health check is done or not.

firstHCDone bool

// isMultiplexed is true if the session is multiplexed.

isMultiplexed bool

}

// isValid returns true if the session is still valid for use.

func (s *session) isValid() bool {

s.mu.Lock()

defer s.mu.Unlock()

return s.valid

}

// isWritePrepared returns true if the session is prepared for write.

func (s *session) isWritePrepared() bool {

s.mu.Lock()

defer s.mu.Unlock()

return s.tx != nil

}

// String implements fmt.Stringer for session.

func (s *session) String() string {

s.mu.Lock()

defer s.mu.Unlock()

return fmt.Sprintf("<id=%v, hcIdx=%v, idleList=%p, valid=%v, create=%v, nextcheck=%v>",

s.id, s.hcIndex, s.idleList, s.valid, s.createTime, s.nextCheck)

}

// ping verifies if the session is still alive in Cloud Spanner.

func (s *session) ping() error {

ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

defer cancel()

// Start parent span that doesn't record.

_, span := octrace.StartSpan(ctx, "cloud.google.com/go/spanner.ping", octrace.WithSampler(octrace.NeverSample()))

defer span.End()

// s.getID is safe even when s is invalid.

_, err := s.client.ExecuteSql(contextWithOutgoingMetadata(ctx, s.md, true), &sppb.ExecuteSqlRequest{

Session: s.getID(),

Sql: "SELECT 1",

})

return err

}

// setHcIndex atomically sets the session's index in the healthcheck queue and

// returns the old index.

func (s *session) setHcIndex(i int) int {

s.mu.Lock()

defer s.mu.Unlock()

oi := s.hcIndex

s.hcIndex = i

return oi

}

// setIdleList atomically sets the session's idle list link and returns the old

// link.

func (s *session) setIdleList(le *list.Element) *list.Element {

s.mu.Lock()

defer s.mu.Unlock()

old := s.idleList

s.idleList = le

return old

}

// invalidate marks a session as invalid and returns the old validity.

func (s *session) invalidate() bool {

s.mu.Lock()

defer s.mu.Unlock()

ov := s.valid

s.valid = false

return ov

}

// setNextCheck sets the timestamp for next healthcheck on the session.

func (s *session) setNextCheck(t time.Time) {

s.mu.Lock()

defer s.mu.Unlock()

s.nextCheck = t

}

// setTransactionID sets the transaction id in the session

func (s *session) setTransactionID(tx transactionID) {

s.mu.Lock()

defer s.mu.Unlock()

s.tx = tx

}

// getID returns the session ID which uniquely identifies the session in Cloud

// Spanner.

func (s *session) getID() string {

s.mu.Lock()

defer s.mu.Unlock()

return s.id

}

// getHcIndex returns the session's index into the global healthcheck priority

// queue.

func (s *session) getHcIndex() int {

s.mu.Lock()

defer s.mu.Unlock()

return s.hcIndex

}

// getIdleList returns the session's link in its home session pool's idle list.

func (s *session) getIdleList() *list.Element {

s.mu.Lock()

defer s.mu.Unlock()

return s.idleList

}

// getNextCheck returns the timestamp for next healthcheck on the session.

func (s *session) getNextCheck() time.Time {

s.mu.Lock()

defer s.mu.Unlock()

return s.nextCheck

}

// recycle turns the session back to its home session pool.

func (s *session) recycle() {

s.setTransactionID(nil)

s.pool.mu.Lock()

if s.isMultiplexed {

s.pool.decNumMultiplexedInUseLocked(context.Background())

s.pool.mu.Unlock()

return

}

if !s.pool.recycleLocked(s) {

// s is rejected by its home session pool because it expired and the

// session pool currently has enough open sessions.

s.pool.mu.Unlock()

// s.destroy internally calls decNumInUseLocked.

s.destroy(false, true)

return

}

s.pool.decNumInUseLocked(context.Background())

s.pool.mu.Unlock()

}

// destroy removes the session from its home session pool, healthcheck queue

// and Cloud Spanner service.

func (s *session) destroy(isExpire, wasInUse bool) bool {

ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)

defer cancel()

return s.destroyWithContext(ctx, isExpire, wasInUse)

}

func (s *session) destroyWithContext(ctx context.Context, isExpire, wasInUse bool) bool {

// Remove s from session pool.

if !s.pool.remove(s, isExpire, wasInUse) {

return false

}

// Unregister s from healthcheck queue.

s.pool.hc.unregister(s)

return true

}

func (s *session) delete(ctx context.Context) {

// Ignore the error because even if we fail to explicitly destroy the

// session, it will be eventually garbage collected by Cloud Spanner.

err := s.client.DeleteSession(contextWithOutgoingMetadata(ctx, s.md, true), &sppb.DeleteSessionRequest{Name: s.getID()})

// Do not log DeadlineExceeded errors when deleting sessions, as these do

// not indicate anything the user can or should act upon.

if err != nil && ErrCode(err) != codes.DeadlineExceeded {

logf(s.logger, "Failed to delete session %v. Error: %v", s.getID(), err)

}

}

// SessionPoolConfig stores configurations of a session pool.

type SessionPoolConfig struct {

// MaxOpened is the maximum number of opened sessions allowed by the session

// pool. If the client tries to open a session and there are already

// MaxOpened sessions, it will block until one becomes available or the

// context passed to the client method is canceled or times out.

//

// Defaults to NumChannels * 100.

MaxOpened uint64

// MinOpened is the minimum number of opened sessions that the session pool

// tries to maintain. Session pool won't continue to expire sessions if

// number of opened connections drops below MinOpened. However, if a session

// is found to be broken, it will still be evicted from the session pool,

// therefore it is posssible that the number of opened sessions drops below

// MinOpened.

//

// Defaults to 100.

MinOpened uint64

// MaxIdle is the maximum number of idle sessions that are allowed in the

// session pool.

//

// Defaults to 0.

MaxIdle uint64

// MaxBurst is the maximum number of concurrent session creation requests.

//

// Deprecated: MaxBurst exists for historical compatibility and should not

// be used. MaxBurst was used to limit the number of sessions that the

// session pool could create within a time frame. This was an early safety

// valve to prevent a client from overwhelming the backend if a large number

// of sessions was suddenly needed. The session pool would then pause the

// creation of sessions for a while. Such a pause is no longer needed and

// the implementation has been removed from the pool.

//

// Defaults to 10.

MaxBurst uint64

// incStep is the number of sessions to create in one batch when at least

// one more session is needed.

//

// Defaults to 25.

incStep uint64

// WriteSessions is the fraction of sessions we try to keep prepared for

// write.

//

// Deprecated: The session pool no longer prepares a fraction of the sessions with a read/write transaction.

// This setting therefore does not have any meaning anymore, and may be removed in the future.

//

// Defaults to 0.2.

WriteSessions float64

// HealthCheckWorkers is number of workers used by health checker for this

// pool.

//

// Defaults to 10.

HealthCheckWorkers int

// HealthCheckInterval is how often the health checker pings a session.

//

// Defaults to 50m.

HealthCheckInterval time.Duration

// MultiplexSessionCheckInterval is the interval at which the multiplexed session is checked whether it needs to be refreshed.

//

// Defaults to 10 mins.

MultiplexSessionCheckInterval time.Duration

// TrackSessionHandles determines whether the session pool will keep track

// of the stacktrace of the goroutines that take sessions from the pool.

// This setting can be used to track down session leak problems.

//

// Defaults to false.

TrackSessionHandles bool

enableMultiplexSession bool

// enableMultiplexedSessionForRW is a flag to enable multiplexed session for read/write transactions

// Defaults to true

enableMultiplexedSessionForRW bool

// enableMultiplexedSessionForPartitionedOps is a flag to enable multiplexed session for partitioned DML and read/query operations

// Defaults to true

enableMultiplexedSessionForPartitionedOps bool

// healthCheckSampleInterval is how often the health checker samples live

// session (for use in maintaining session pool size).

//

// Defaults to 1m.

healthCheckSampleInterval time.Duration

// sessionLabels for the sessions created in the session pool.

sessionLabels map[string]string

InactiveTransactionRemovalOptions

}

// DefaultSessionPoolConfig is the default configuration for the session pool

// that will be used for a Spanner client, unless the user supplies a specific

// session pool config.

var DefaultSessionPoolConfig = SessionPoolConfig{

MinOpened: 100,

MaxOpened: numChannels * 100,

MaxBurst: 10,

incStep: 25,

WriteSessions: 0.2,

HealthCheckWorkers: 10,

HealthCheckInterval: healthCheckIntervalMins * time.Minute,

InactiveTransactionRemovalOptions: InactiveTransactionRemovalOptions{

ActionOnInactiveTransaction: Warn,

executionFrequency: 2 * time.Minute,

idleTimeThreshold: 60 * time.Minute,

usedSessionsRatioThreshold: 0.95,

},

}

// errMinOpenedGTMapOpened returns error for SessionPoolConfig.MaxOpened < SessionPoolConfig.MinOpened when SessionPoolConfig.MaxOpened is set.

func errMinOpenedGTMaxOpened(maxOpened, minOpened uint64) error {

return spannerErrorf(codes.InvalidArgument,

"require SessionPoolConfig.MaxOpened >= SessionPoolConfig.MinOpened, got %d and %d", maxOpened, minOpened)

}

// errWriteFractionOutOfRange returns error for

// SessionPoolConfig.WriteFraction < 0 or SessionPoolConfig.WriteFraction > 1

func errWriteFractionOutOfRange(writeFraction float64) error {

return spannerErrorf(codes.InvalidArgument,

"require SessionPoolConfig.WriteSessions >= 0.0 && SessionPoolConfig.WriteSessions <= 1.0, got %.2f", writeFraction)

}

// errHealthCheckWorkersNegative returns error for

// SessionPoolConfig.HealthCheckWorkers < 0

func errHealthCheckWorkersNegative(workers int) error {

return spannerErrorf(codes.InvalidArgument,

"require SessionPoolConfig.HealthCheckWorkers >= 0, got %d", workers)

}

// errHealthCheckIntervalNegative returns error for

// SessionPoolConfig.HealthCheckInterval < 0

func errHealthCheckIntervalNegative(interval time.Duration) error {

return spannerErrorf(codes.InvalidArgument,

"require SessionPoolConfig.HealthCheckInterval >= 0, got %v", interval)

}

// validate verifies that the SessionPoolConfig is good for use.

func (spc *SessionPoolConfig) validate() error {

if spc.MinOpened > spc.MaxOpened && spc.MaxOpened > 0 {

return errMinOpenedGTMaxOpened(spc.MaxOpened, spc.MinOpened)

}

if spc.HealthCheckWorkers < 0 {

return errHealthCheckWorkersNegative(spc.HealthCheckWorkers)

}

if spc.HealthCheckInterval < 0 {

return errHealthCheckIntervalNegative(spc.HealthCheckInterval)

}

return nil

}

type muxSessionCreateRequest struct {

ctx context.Context

force bool

}

// sessionPool creates and caches Cloud Spanner sessions.

type sessionPool struct {

// mu protects sessionPool from concurrent access.

mu sync.Mutex

// valid marks the validity of the session pool.

valid bool

// sc is used to create the sessions for the pool.

sc *sessionClient

// trackedSessionHandles contains all sessions handles that have been

// checked out of the pool. The list is only filled if TrackSessionHandles

// has been enabled.

trackedSessionHandles list.List

// idleList caches idle session IDs. Session IDs in this list can be

// allocated for use.

idleList list.List

// multiplexSessionClientCounter is the counter for the multiplexed session client.

multiplexSessionClientCounter int

// clientPool is a pool of Cloud Spanner grpc clients.

clientPool []spannerClient

// multiplexedSession contains the multiplexed session

multiplexedSession *session

// mayGetSession is for broadcasting that session retrival/creation may

// proceed.

mayGetSession chan struct{}

// multiplexedSessionReq is the ongoing multiplexed session creation request (if any).

multiplexedSessionReq chan muxSessionCreateRequest

// mayGetMultiplexedSession is for broadcasting that multiplexed session retrieval is possible.

mayGetMultiplexedSession chan bool

// sessionCreationError is the last error that occurred during session

// creation and is propagated to any waiters waiting for a session.

sessionCreationError error

// multiplexedSessionCreationError is the error that occurred during multiplexed session

// creation for the first time and is propagated to any waiters waiting for a session.

multiplexedSessionCreationError error

// numOpened is the total number of open sessions from the session pool.

numOpened uint64

// createReqs is the number of ongoing session creation requests.

createReqs uint64

// numWaiters is the number of processes waiting for a session to

// become available.

numWaiters uint64

// disableBackgroundPrepareSessions indicates that the BeginTransaction

// call for a read/write transaction failed with a permanent error, such as

// PermissionDenied or `Database not found`. Further background calls to

// prepare sessions will be disabled.

disableBackgroundPrepareSessions bool

// configuration of the session pool.

SessionPoolConfig

// hc is the health checker

hc *healthChecker

// rand is a separately sourced random generator.

rand *rand.Rand

// numInUse is the number of sessions that are currently in use (checked out

// from the session pool).

numInUse uint64

// maxNumInUse is the maximum number of sessions in use concurrently in the

// current 10 minute interval.

maxNumInUse uint64

// lastResetTime is the start time of the window for recording maxNumInUse.

lastResetTime time.Time

// numSessions is the number of sessions that are idle for read/write.

numSessions uint64

// mw is the maintenance window containing statistics for the max number of

// sessions checked out of the pool during the last 10 minutes.

mw *maintenanceWindow

// tagMap is a map of all tags that are associated with the emitted metrics.

tagMap *tag.Map

// indicates the number of leaked sessions removed from the session pool.

// This is valid only when ActionOnInactiveTransaction is WarnAndClose or ActionOnInactiveTransaction is Close in InactiveTransactionRemovalOptions.

numOfLeakedSessionsRemoved uint64

otConfig *openTelemetryConfig

// enableMultiplexSession is a flag to enable multiplexed session.

enableMultiplexSession bool

}

// newSessionPool creates a new session pool.

func newSessionPool(sc *sessionClient, config SessionPoolConfig) (*sessionPool, error) {

if err := config.validate(); err != nil {

return nil, err

}

if config.HealthCheckWorkers == 0 {

// With 10 workers and assuming average latency of 5ms for

// BeginTransaction, we will be able to prepare 2000 tx/sec in advance.

// If the rate of takeWriteSession is more than that, it will degrade to

// doing BeginTransaction inline.

//

// TODO: consider resizing the worker pool dynamically according to the load.

config.HealthCheckWorkers = 10

}

if config.HealthCheckInterval == 0 {

config.HealthCheckInterval = healthCheckIntervalMins * time.Minute

}

if config.healthCheckSampleInterval == 0 {

config.healthCheckSampleInterval = time.Minute

}

if config.ActionOnInactiveTransaction == actionUnspecified {

config.ActionOnInactiveTransaction = DefaultSessionPoolConfig.ActionOnInactiveTransaction

}

if config.idleTimeThreshold == 0 {

config.idleTimeThreshold = DefaultSessionPoolConfig.idleTimeThreshold

}

if config.executionFrequency == 0 {

config.executionFrequency = DefaultSessionPoolConfig.executionFrequency

}

if config.usedSessionsRatioThreshold == 0 {

config.usedSessionsRatioThreshold = DefaultSessionPoolConfig.usedSessionsRatioThreshold

}

if config.MultiplexSessionCheckInterval == 0 {

config.MultiplexSessionCheckInterval = 10 * time.Minute

}

pool := &sessionPool{

sc: sc,

valid: true,

mayGetSession: make(chan struct{}),

mayGetMultiplexedSession: make(chan bool),

multiplexedSessionReq: make(chan muxSessionCreateRequest),

SessionPoolConfig: config,

mw: newMaintenanceWindow(config.MaxOpened),

rand: rand.New(rand.NewSource(time.Now().UnixNano())),

otConfig: sc.otConfig,

enableMultiplexSession: config.enableMultiplexSession,

}

_, instance, database, err := parseDatabaseName(sc.database)

if err != nil {

return nil, err

}

// Errors should not prevent initializing the session pool.

ctx, err := tag.New(context.Background(),

tag.Upsert(tagKeyClientID, sc.id),

tag.Upsert(tagKeyDatabase, database),

tag.Upsert(tagKeyInstance, instance),

tag.Upsert(tagKeyLibVersion, internal.Version),

)

if err != nil {

logf(pool.sc.logger, "Failed to create tag map, error: %v", err)

}

pool.tagMap = tag.FromContext(ctx)

// On GCE VM, within the same region an healthcheck ping takes on average

// 10ms to finish, given a 5 minutes interval and 10 healthcheck workers, a

// healthChecker can effectively mantain

// 100 checks_per_worker/sec * 10 workers * 300 seconds = 300K sessions.

pool.hc = newHealthChecker(config.HealthCheckInterval, config.MultiplexSessionCheckInterval, config.HealthCheckWorkers, config.healthCheckSampleInterval, pool)

// First initialize the pool before we indicate that the healthchecker is

// ready. This prevents the maintainer from starting before the pool has

// been initialized, which means that we guarantee that the initial

// sessions are created using BatchCreateSessions.

if config.MinOpened > 0 {

numSessions := minUint64(config.MinOpened, math.MaxInt32)

if err := pool.initPool(numSessions); err != nil {

return nil, err

}

}

if pool.enableMultiplexSession {

go pool.createMultiplexedSession()

ctx, cancel := context.WithTimeout(ctx, 30*time.Second)

pool.multiplexedSessionReq <- muxSessionCreateRequest{force: true, ctx: ctx}

// listen for the session to be created

go func() {

select {

case <-ctx.Done():

cancel()

return

// wait for the session to be created

case <-pool.mayGetMultiplexedSession:

}

}()

}

pool.recordStat(context.Background(), MaxAllowedSessionsCount, int64(config.MaxOpened))

err = registerSessionPoolOTMetrics(pool)

if err != nil {

logf(pool.sc.logger, "Error when registering session pool metrics in OpenTelemetry, error: %v", err)

}

close(pool.hc.ready)

return pool, nil

}

func (p *sessionPool) recordStat(ctx context.Context, m *stats.Int64Measure, n int64, tags ...tag.Tag) {

ctx = tag.NewContext(ctx, p.tagMap)

mutators := make([]tag.Mutator, len(tags))

for i, t := range tags {

mutators[i] = tag.Upsert(t.Key, t.Value)

}

ctx, err := tag.New(ctx, mutators...)

if err != nil {

logf(p.sc.logger, "Failed to tag metrics, error: %v", err)

}

recordStat(ctx, m, n)

}

type recordOTStatOption struct {

attr []attribute.KeyValue

}

func (p *sessionPool) recordOTStat(ctx context.Context, m metric.Int64Counter, val int64, option recordOTStatOption) {

if m != nil {

attrs := p.otConfig.attributeMap

if len(option.attr) > 0 {

attrs = option.attr

}

m.Add(ctx, val, metric.WithAttributes(attrs...))

}

}

func (p *sessionPool) getRatioOfSessionsInUseLocked() float64 {

maxSessions := p.MaxOpened

if maxSessions == 0 {

return 0

}

return float64(p.numInUse) / float64(maxSessions)

}

func (p *sessionPool) isMultiplexedSessionForRWEnabled() bool {

p.mu.Lock()

defer p.mu.Unlock()

return p.enableMultiplexedSessionForRW

}

func (p *sessionPool) isMultiplexedSessionForPartitionedOpsEnabled() bool {

p.mu.Lock()

defer p.mu.Unlock()

return p.enableMultiplexedSessionForPartitionedOps

}

func (p *sessionPool) disableMultiplexedSessionForRW() {

p.mu.Lock()

defer p.mu.Unlock()

p.enableMultiplexedSessionForRW = false

}

func (p *sessionPool) disableMultiplexedSessionForPartitionedOps() {

p.mu.Lock()

defer p.mu.Unlock()

p.enableMultiplexedSessionForPartitionedOps = false

}

// gets sessions which are unexpectedly long-running.

func (p *sessionPool) getLongRunningSessionsLocked() []*sessionHandle {

usedSessionsRatio := p.getRatioOfSessionsInUseLocked()

var longRunningSessions []*sessionHandle

if usedSessionsRatio > p.usedSessionsRatioThreshold {

element := p.trackedSessionHandles.Front()

for element != nil {

sh := element.Value.(*sessionHandle)

sh.mu.Lock()

if sh.session == nil {

// sessionHandle has already been recycled/destroyed.

sh.mu.Unlock()

element = element.Next()

continue

}

diff := time.Since(sh.lastUseTime)

if !sh.eligibleForLongRunning && diff.Seconds() >= p.idleTimeThreshold.Seconds() {

if (p.ActionOnInactiveTransaction == Warn || p.ActionOnInactiveTransaction == WarnAndClose) && !sh.isSessionLeakLogged {

if p.ActionOnInactiveTransaction == Warn {

if sh.stack != nil {

logf(p.sc.logger, "session %s checked out of pool at %s is long running due to possible session leak for goroutine: \n%s", sh.session.getID(), sh.checkoutTime.Format(time.RFC3339), sh.stack)

} else {

logf(p.sc.logger, "session %s checked out of pool at %s is long running due to possible session leak for goroutine: \nEnable SessionPoolConfig.TrackSessionHandles to get stack trace associated with the session", sh.session.getID(), sh.checkoutTime.Format(time.RFC3339))

}

sh.isSessionLeakLogged = true

} else if p.ActionOnInactiveTransaction == WarnAndClose {

if sh.stack != nil {

logf(p.sc.logger, "session %s checked out of pool at %s is long running and will be removed due to possible session leak for goroutine: \n%s", sh.session.getID(), sh.checkoutTime.Format(time.RFC3339), sh.stack)

} else {

logf(p.sc.logger, "session %s checked out of pool at %s is long running and will be removed due to possible session leak for goroutine: \nEnable SessionPoolConfig.TrackSessionHandles to get stack trace associated with the session", sh.session.getID(), sh.checkoutTime.Format(time.RFC3339))

}

}

}

if p.ActionOnInactiveTransaction == WarnAndClose || p.ActionOnInactiveTransaction == Close {

longRunningSessions = append(longRunningSessions, sh)

}

}

sh.mu.Unlock()

element = element.Next()

}

}

return longRunningSessions

}

// removes or logs sessions that are unexpectedly long-running.

func (p *sessionPool) removeLongRunningSessions() {

p.mu.Lock()

longRunningSessions := p.getLongRunningSessionsLocked()

p.mu.Unlock()

// destroy long-running sessions

if p.ActionOnInactiveTransaction == WarnAndClose || p.ActionOnInactiveTransaction == Close {

var leakedSessionsRemovedCount uint64

for _, sh := range longRunningSessions {

// removes inner session out of the pool to reduce the probability of two processes trying

// to use the same session at the same time.

sh.destroy()

leakedSessionsRemovedCount++

}

p.mu.Lock()

p.numOfLeakedSessionsRemoved += leakedSessionsRemovedCount

p.mu.Unlock()

}

}

func (p *sessionPool) initPool(numSessions uint64) error {

p.mu.Lock()

defer p.mu.Unlock()

return p.growPoolLocked(numSessions, true)

}

func (p *sessionPool) growPoolLocked(numSessions uint64, distributeOverChannels bool) error {

// Take budget before the actual session creation.

numSessions = minUint64(numSessions, math.MaxInt32)

p.numOpened += uint64(numSessions)

p.recordStat(context.Background(), OpenSessionCount, int64(p.numOpened))

p.createReqs += uint64(numSessions)

// Asynchronously create a batch of sessions for the pool.

return p.sc.batchCreateSessions(int32(numSessions), distributeOverChannels, p)

}

func (p *sessionPool) createMultiplexedSession() {

for c := range p.multiplexedSessionReq {

p.mu.Lock()

sess := p.multiplexedSession

p.mu.Unlock()

if c.force || sess == nil {

p.mu.Lock()

p.sc.mu.Lock()

client, err := p.sc.nextClient()

p.sc.mu.Unlock()

p.mu.Unlock()

if err != nil {

// If we can't get a client, we can't create a session.

p.mu.Lock()

p.multiplexedSessionCreationError = err

p.mu.Unlock()

select {

case p.mayGetMultiplexedSession <- true:

case <-c.ctx.Done():

}

continue

}

p.sc.executeCreateMultiplexedSession(c.ctx, client, p.sc.md, p)

continue

}

select {

case p.mayGetMultiplexedSession <- true:

case <-c.ctx.Done():

return

}

}

}

// sessionReady is executed by the SessionClient when a session has been

// created and is ready to use. This method will add the new session to the

// pool and decrease the number of sessions that is being created.

func (p *sessionPool) sessionReady(ctx context.Context, s *session) {

p.mu.Lock()

defer p.mu.Unlock()

// Clear any session creation error.

if s.isMultiplexed {

s.pool = p

p.multiplexedSession = s

p.multiplexedSessionCreationError = nil

p.recordStat(context.Background(), OpenSessionCount, int64(1), tag.Tag{Key: tagKeyIsMultiplexed, Value: "true"})

p.recordStat(context.Background(), SessionsCount, 1, tagNumSessions, tag.Tag{Key: tagKeyIsMultiplexed, Value: "true"})

// either notify the waiting goroutine or skip if no one is waiting

select {

case p.mayGetMultiplexedSession <- true:

case <-ctx.Done():

return

}

return

}

p.sessionCreationError = nil

// Set this pool as the home pool of the session and register it with the

// health checker.

s.pool = p

p.hc.register(s)

p.createReqs--

// Insert the session at a random position in the pool to prevent all

// sessions affiliated with a channel to be placed at sequentially in the

// pool.

if p.idleList.Len() > 0 {

pos := rand.Intn(p.idleList.Len())

before := p.idleList.Front()

for i := 0; i < pos; i++ {

before = before.Next()

}

s.setIdleList(p.idleList.InsertBefore(s, before))

} else {

s.setIdleList(p.idleList.PushBack(s))

}

p.incNumSessionsLocked(context.Background())

// Notify other waiters blocking on session creation.

close(p.mayGetSession)

p.mayGetSession = make(chan struct{})