// getConn reuses an existing connection if possible. Otherwise
// it returns a connection which it will save for future reuse.
// If it returns an error, retry will tell you if getConn can be retried.
// If the context has a deadline and exceeded, it returns error and no-retry immediately.
func (sdc *ShardConn) getConn(ctx context.Context) (conn tabletconn.TabletConn, endPoint topo.EndPoint, err error, retry bool) {
sdc.mu.Lock()
defer sdc.mu.Unlock()
// fail-fast if deadline exceeded
deadline, ok := ctx.Deadline()
if ok {
if time.Now().After(deadline) {
return nil, topo.EndPoint{}, tabletconn.OperationalError("vttablet: deadline exceeded"), false
}
}
if sdc.conn != nil {
return sdc.conn, sdc.conn.EndPoint(), nil, false
}
endPoint, err = sdc.balancer.Get()
if err != nil {
return nil, topo.EndPoint{}, err, false
}
conn, err = tabletconn.GetDialer()(ctx, endPoint, sdc.keyspace, sdc.shard, sdc.timeout)
if err != nil {
sdc.balancer.MarkDown(endPoint.Uid, err.Error())
return nil, endPoint, err, true
}
sdc.conn = conn
return sdc.conn, endPoint, nil, false
}
func (sbc *sandboxConn) getError() error {
if sbc.onConnUse != nil {
sbc.onConnUse(sbc)
}
if sbc.mustFailRetry > 0 {
sbc.mustFailRetry--
return &tabletconn.ServerError{Code: tabletconn.ERR_RETRY, Err: "retry: err"}
}
if sbc.mustFailFatal > 0 {
sbc.mustFailFatal--
return &tabletconn.ServerError{Code: tabletconn.ERR_FATAL, Err: "fatal: err"}
}
if sbc.mustFailServer > 0 {
sbc.mustFailServer--
return &tabletconn.ServerError{Code: tabletconn.ERR_NORMAL, Err: "error: err"}
}
if sbc.mustFailConn > 0 {
sbc.mustFailConn--
return tabletconn.OperationalError(fmt.Sprintf("error: conn"))
}
if sbc.mustFailTxPool > 0 {
sbc.mustFailTxPool--
return &tabletconn.ServerError{Code: tabletconn.ERR_TX_POOL_FULL, Err: "tx_pool_full: err"}
}
if sbc.mustFailNotTx > 0 {
sbc.mustFailNotTx--
return &tabletconn.ServerError{Code: tabletconn.ERR_NOT_IN_TX, Err: "not_in_tx: err"}
}
return nil
}
func sandboxDialer(context context.Context, endPoint topo.EndPoint, keyspace, shard string, timeout time.Duration) (tabletconn.TabletConn, error) {
sand := getSandbox(keyspace)
sand.sandmu.Lock()
defer sand.sandmu.Unlock()
sand.DialCounter++
if sand.DialMustFail > 0 {
sand.DialMustFail--
return nil, tabletconn.OperationalError(fmt.Sprintf("conn error"))
}
conns := sand.TestConns[shard]
if conns == nil {
panic(fmt.Sprintf("can't find shard %v", shard))
}
tconn := conns[endPoint.Uid]
tconn.(*sandboxConn).endPoint = endPoint
return tconn, nil
}
// withRetry sets up the connection and executes the action. If there are connection errors,
// it retries retryCount times before failing. It does not retry if the connection is in
// the middle of a transaction. While returning the error check if it maybe a result of
// a resharding event, and set the re-resolve bit and let the upper layers
// re-resolve and retry.
func (sdc *ShardConn) withRetry(ctx context.Context, action func(conn tabletconn.TabletConn) error, transactionID int64, isStreaming bool) error {
var conn tabletconn.TabletConn
var endPoint topo.EndPoint
var err error
var retry bool
inTransaction := (transactionID != 0)
// execute the action at least once even without retrying
for i := 0; i < sdc.retryCount+1; i++ {
conn, endPoint, err, retry = sdc.getConn(ctx)
if err != nil {
if retry {
continue
}
return sdc.WrapError(err, endPoint, inTransaction)
}
// no timeout for streaming query
if isStreaming {
err = action(conn)
} else {
tmr := time.NewTimer(sdc.timeout)
done := make(chan int)
var errAction error
go func() {
errAction = action(conn)
close(done)
}()
select {
case <-tmr.C:
err = tabletconn.OperationalError("vttablet: call timeout")
case <-done:
err = errAction
}
tmr.Stop()
}
if sdc.canRetry(err, transactionID, conn) {
continue
}
return sdc.WrapError(err, endPoint, inTransaction)
}
return sdc.WrapError(err, endPoint, inTransaction)
}
func tabletError(err error) error {
if err == nil {
return nil
}
if _, ok := err.(rpcplus.ServerError); ok {
var code int
errStr := err.Error()
switch {
case strings.HasPrefix(errStr, "fatal"):
code = tabletconn.ERR_FATAL
case strings.HasPrefix(errStr, "retry"):
code = tabletconn.ERR_RETRY
case strings.HasPrefix(errStr, "tx_pool_full"):
code = tabletconn.ERR_TX_POOL_FULL
case strings.HasPrefix(errStr, "not_in_tx"):
code = tabletconn.ERR_NOT_IN_TX
default:
code = tabletconn.ERR_NORMAL
}
return &tabletconn.ServerError{Code: code, Err: fmt.Sprintf("vttablet: %v", err)}
}
return tabletconn.OperationalError(fmt.Sprintf("vttablet: %v", err))
}