vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletserver/stateful_connection.go (about)

     1  /*
     2  Copyright 2019 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  	"fmt"
    22  	"time"
    23  
    24  	"vitess.io/vitess/go/mysql"
    25  	"vitess.io/vitess/go/pools"
    26  	"vitess.io/vitess/go/sqltypes"
    27  	"vitess.io/vitess/go/vt/callerid"
    28  	"vitess.io/vitess/go/vt/log"
    29  	"vitess.io/vitess/go/vt/servenv"
    30  	"vitess.io/vitess/go/vt/vterrors"
    31  	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"
    32  	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
    33  	"vitess.io/vitess/go/vt/vttablet/tabletserver/tx"
    34  
    35  	querypb "vitess.io/vitess/go/vt/proto/query"
    36  	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
    37  )
    38  
    39  // StatefulConnection is used in the situations where we need a dedicated connection for a vtgate session.
    40  // This is used for transactions and reserved connections.
    41  // NOTE: After use, if must be returned either by doing a Unlock() or a Release().
    42  type StatefulConnection struct {
    43  	pool           *StatefulConnectionPool
    44  	dbConn         *connpool.DBConn
    45  	ConnID         tx.ConnID
    46  	env            tabletenv.Env
    47  	txProps        *tx.Properties
    48  	reservedProps  *Properties
    49  	tainted        bool
    50  	enforceTimeout bool
    51  	timeout        time.Duration
    52  	expiryTime     time.Time
    53  }
    54  
    55  // Properties contains meta information about the connection
    56  type Properties struct {
    57  	EffectiveCaller *vtrpcpb.CallerID
    58  	ImmediateCaller *querypb.VTGateCallerID
    59  	StartTime       time.Time
    60  	Stats           *servenv.TimingsWrapper
    61  }
    62  
    63  // Close closes the underlying connection. When the connection is Unblocked, it will be Released
    64  func (sc *StatefulConnection) Close() {
    65  	if sc.dbConn != nil {
    66  		sc.dbConn.Close()
    67  	}
    68  }
    69  
    70  // IsClosed returns true when the connection is still operational
    71  func (sc *StatefulConnection) IsClosed() bool {
    72  	return sc.dbConn == nil || sc.dbConn.IsClosed()
    73  }
    74  
    75  // IsInTransaction returns true when the connection has tx state
    76  func (sc *StatefulConnection) IsInTransaction() bool {
    77  	return sc.txProps != nil
    78  }
    79  
    80  func (sc *StatefulConnection) ElapsedTimeout() bool {
    81  	if !sc.enforceTimeout {
    82  		return false
    83  	}
    84  	if sc.timeout <= 0 {
    85  		return false
    86  	}
    87  	return sc.expiryTime.Before(time.Now())
    88  }
    89  
    90  // Exec executes the statement in the dedicated connection
    91  func (sc *StatefulConnection) Exec(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
    92  	if sc.IsClosed() {
    93  		if sc.IsInTransaction() {
    94  			return nil, vterrors.Errorf(vtrpcpb.Code_ABORTED, "transaction was aborted: %v", sc.txProps.Conclusion)
    95  		}
    96  		return nil, vterrors.New(vtrpcpb.Code_ABORTED, "connection was aborted")
    97  	}
    98  	r, err := sc.dbConn.ExecOnce(ctx, query, maxrows, wantfields)
    99  	if err != nil {
   100  		if mysql.IsConnErr(err) {
   101  			select {
   102  			case <-ctx.Done():
   103  				// If the context is done, the query was killed.
   104  				// So, don't trigger a mysql check.
   105  			default:
   106  				sc.env.CheckMySQL()
   107  			}
   108  			return nil, err
   109  		}
   110  		return nil, err
   111  	}
   112  	return r, nil
   113  }
   114  
   115  func (sc *StatefulConnection) execWithRetry(ctx context.Context, query string, maxrows int, wantfields bool) (string, error) {
   116  	if sc.IsClosed() {
   117  		return "", vterrors.New(vtrpcpb.Code_CANCELED, "connection is closed")
   118  	}
   119  	res, err := sc.dbConn.Exec(ctx, query, maxrows, wantfields)
   120  	if err != nil {
   121  		return "", err
   122  	}
   123  	return res.SessionStateChanges, nil
   124  }
   125  
   126  // FetchNext returns the next result set.
   127  func (sc *StatefulConnection) FetchNext(ctx context.Context, maxrows int, wantfields bool) (*sqltypes.Result, error) {
   128  	if sc.IsClosed() {
   129  		return nil, vterrors.New(vtrpcpb.Code_CANCELED, "connection is closed")
   130  	}
   131  	return sc.dbConn.FetchNext(ctx, maxrows, wantfields)
   132  }
   133  
   134  // Unlock returns the connection to the pool. The connection remains active.
   135  // This method is idempotent and can be called multiple times
   136  func (sc *StatefulConnection) Unlock() {
   137  	// when in a transaction, we count from the time created, so each use of the connection does not update the time
   138  	updateTime := !sc.IsInTransaction()
   139  	sc.unlock(updateTime)
   140  }
   141  
   142  // UnlockUpdateTime returns the connection to the pool. The connection remains active.
   143  // This method is idempotent and can be called multiple times
   144  func (sc *StatefulConnection) UnlockUpdateTime() {
   145  	sc.unlock(true)
   146  }
   147  
   148  func (sc *StatefulConnection) unlock(updateTime bool) {
   149  	if sc.dbConn == nil {
   150  		return
   151  	}
   152  	if sc.dbConn.IsClosed() {
   153  		sc.Releasef("unlocked closed connection")
   154  	} else {
   155  		sc.pool.markAsNotInUse(sc, updateTime)
   156  	}
   157  }
   158  
   159  // Release is used when the connection will not be used ever again.
   160  // The underlying dbConn is removed so that this connection cannot be used by mistake.
   161  func (sc *StatefulConnection) Release(reason tx.ReleaseReason) {
   162  	sc.Releasef(reason.String())
   163  }
   164  
   165  // Releasef is used when the connection will not be used ever again.
   166  // The underlying dbConn is removed so that this connection cannot be used by mistake.
   167  func (sc *StatefulConnection) Releasef(reasonFormat string, a ...any) {
   168  	if sc.dbConn == nil {
   169  		return
   170  	}
   171  	sc.pool.unregister(sc.ConnID, fmt.Sprintf(reasonFormat, a...))
   172  	sc.dbConn.Recycle()
   173  	sc.dbConn = nil
   174  	sc.logReservedConn()
   175  }
   176  
   177  // Renew the existing connection with new connection id.
   178  func (sc *StatefulConnection) Renew() error {
   179  	err := sc.pool.renewConn(sc)
   180  	if err != nil {
   181  		sc.Close()
   182  		return vterrors.Wrap(err, "connection renew failed")
   183  	}
   184  	return nil
   185  }
   186  
   187  // String returns a printable version of the connection info.
   188  func (sc *StatefulConnection) String(sanitize bool) string {
   189  	return fmt.Sprintf(
   190  		"%v\t%s",
   191  		sc.ConnID,
   192  		sc.txProps.String(sanitize),
   193  	)
   194  }
   195  
   196  // Current returns the currently executing query
   197  func (sc *StatefulConnection) Current() string {
   198  	return sc.dbConn.Current()
   199  }
   200  
   201  // ID returns the mysql connection ID
   202  func (sc *StatefulConnection) ID() int64 {
   203  	return sc.dbConn.ID()
   204  }
   205  
   206  // Kill kills the currently executing query and connection
   207  func (sc *StatefulConnection) Kill(reason string, elapsed time.Duration) error {
   208  	return sc.dbConn.Kill(reason, elapsed)
   209  }
   210  
   211  // TxProperties returns the transactional properties of the connection
   212  func (sc *StatefulConnection) TxProperties() *tx.Properties {
   213  	return sc.txProps
   214  }
   215  
   216  // ReservedID returns the identifier for this connection
   217  func (sc *StatefulConnection) ReservedID() tx.ConnID {
   218  	return sc.ConnID
   219  }
   220  
   221  // UnderlyingDBConn returns the underlying database connection
   222  func (sc *StatefulConnection) UnderlyingDBConn() *connpool.DBConn {
   223  	return sc.dbConn
   224  }
   225  
   226  // CleanTxState cleans out the current transaction state
   227  func (sc *StatefulConnection) CleanTxState() {
   228  	sc.txProps = nil
   229  }
   230  
   231  // Stats implements the tx.IStatefulConnection interface
   232  func (sc *StatefulConnection) Stats() *tabletenv.Stats {
   233  	return sc.env.Stats()
   234  }
   235  
   236  // Taint taints the existing connection.
   237  func (sc *StatefulConnection) Taint(ctx context.Context, stats *servenv.TimingsWrapper) error {
   238  	if sc.dbConn == nil {
   239  		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "connection is closed")
   240  	}
   241  	if sc.tainted {
   242  		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "connection is already reserved")
   243  	}
   244  	immediateCaller := callerid.ImmediateCallerIDFromContext(ctx)
   245  	effectiveCaller := callerid.EffectiveCallerIDFromContext(ctx)
   246  
   247  	sc.tainted = true
   248  	sc.reservedProps = &Properties{
   249  		EffectiveCaller: effectiveCaller,
   250  		ImmediateCaller: immediateCaller,
   251  		StartTime:       time.Now(),
   252  		Stats:           stats,
   253  	}
   254  	sc.dbConn.Taint()
   255  	sc.Stats().UserActiveReservedCount.Add(sc.getUsername(), 1)
   256  	return nil
   257  }
   258  
   259  // IsTainted tells us whether this connection is tainted
   260  func (sc *StatefulConnection) IsTainted() bool {
   261  	return sc.tainted
   262  }
   263  
   264  // LogTransaction logs transaction related stats
   265  func (sc *StatefulConnection) LogTransaction(reason tx.ReleaseReason) {
   266  	if sc.txProps == nil {
   267  		return //Nothing to log as no transaction exists on this connection.
   268  	}
   269  	sc.txProps.Conclusion = reason.Name()
   270  	sc.txProps.EndTime = time.Now()
   271  
   272  	username := callerid.GetPrincipal(sc.txProps.EffectiveCaller)
   273  	if username == "" {
   274  		username = callerid.GetUsername(sc.txProps.ImmediateCaller)
   275  	}
   276  	duration := sc.txProps.EndTime.Sub(sc.txProps.StartTime)
   277  	sc.Stats().UserTransactionCount.Add([]string{username, reason.Name()}, 1)
   278  	sc.Stats().UserTransactionTimesNs.Add([]string{username, reason.Name()}, int64(duration))
   279  	sc.txProps.Stats.Add(reason.Name(), duration)
   280  	if sc.txProps.LogToFile {
   281  		log.Infof("Logged transaction: %s", sc.String(sc.env.Config().SanitizeLogMessages))
   282  	}
   283  	tabletenv.TxLogger.Send(sc)
   284  }
   285  
   286  func (sc *StatefulConnection) SetTimeout(timeout time.Duration) {
   287  	sc.timeout = timeout
   288  	sc.resetExpiryTime()
   289  }
   290  
   291  // logReservedConn logs reserved connection related stats.
   292  func (sc *StatefulConnection) logReservedConn() {
   293  	if sc.reservedProps == nil {
   294  		return //Nothing to log as this connection is not reserved.
   295  	}
   296  	duration := time.Since(sc.reservedProps.StartTime)
   297  	username := sc.getUsername()
   298  	sc.Stats().UserActiveReservedCount.Add(username, -1)
   299  	sc.Stats().UserReservedCount.Add(username, 1)
   300  	sc.Stats().UserReservedTimesNs.Add(username, int64(duration))
   301  }
   302  
   303  func (sc *StatefulConnection) getUsername() string {
   304  	username := callerid.GetPrincipal(sc.reservedProps.EffectiveCaller)
   305  	if username != "" {
   306  		return username
   307  	}
   308  	return callerid.GetUsername(sc.reservedProps.ImmediateCaller)
   309  }
   310  
   311  func (sc *StatefulConnection) ApplySetting(ctx context.Context, setting *pools.Setting) error {
   312  	if sc.dbConn.IsSameSetting(setting.GetQuery()) {
   313  		return nil
   314  	}
   315  	return sc.dbConn.ApplySetting(ctx, setting)
   316  }
   317  
   318  func (sc *StatefulConnection) resetExpiryTime() {
   319  	sc.expiryTime = time.Now().Add(sc.timeout)
   320  }