func (s *sandbox) DeleteTestConn(shard string, conn tabletconn.TabletConn) {
s.sandmu.Lock()
defer s.sandmu.Unlock()
conns, ok := s.TestConns[shard]
if !ok {
panic(fmt.Sprintf("unknown shard: %v", shard))
}
delete(conns, conn.EndPoint().Uid)
s.TestConns[shard] = conns
}
// WrapError adds the connection context to an error.
func (sdc *ShardConn) WrapError(in error, conn tabletconn.TabletConn) (wrapped error) {
if in == nil {
return nil
}
var endPoint topo.EndPoint
if conn != nil {
endPoint = conn.EndPoint()
}
return fmt.Errorf(
"%v, shard: (%s.%s.%s), host: %s",
in, sdc.keyspace, sdc.shard, sdc.tabletType, endPoint.Host)
}
// markDown closes conn and temporarily marks the associated
// end point as unusable.
func (sdc *ShardConn) markDown(conn tabletconn.TabletConn) {
sdc.mu.Lock()
defer sdc.mu.Unlock()
if conn != sdc.conn {
return
}
sdc.balancer.MarkDown(conn.EndPoint().Uid)
// Launch as goroutine so we don't block
go sdc.conn.Close()
sdc.conn = nil
}
// markDown closes conn and temporarily marks the associated
// end point as unusable.
func (sdc *ShardConn) markDown(conn tabletconn.TabletConn, reason string) {
sdc.mu.Lock()
defer sdc.mu.Unlock()
if conn != sdc.conn {
return
}
sdc.balancer.MarkDown(conn.EndPoint().Uid, reason)
go func(conn tabletconn.TabletConn) {
danglingTabletConn.Add(1)
conn.Close()
danglingTabletConn.Add(-1)
}(sdc.conn)
sdc.conn = nil
}
// StreamExecute executes a streaming query on vttablet. The retry rules are the same as Execute.
func (sdc *ShardConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (<-chan *mproto.QueryResult, tabletconn.ErrFunc) {
var usedConn tabletconn.TabletConn
var erFunc tabletconn.ErrFunc
var results <-chan *mproto.QueryResult
err := sdc.withRetry(ctx, func(conn tabletconn.TabletConn) error {
results, erFunc = conn.StreamExecute(ctx, query, bindVars, transactionID)
usedConn = conn
return erFunc()
}, transactionID, true)
if err != nil {
return results, func() error { return err }
}
inTransaction := (transactionID != 0)
return results, func() error { return sdc.WrapError(erFunc(), usedConn.EndPoint(), inTransaction) }
}
// StreamExecute executes a streaming query for the specified keyspace, shard, and tablet type.
func (dg *discoveryGateway) StreamExecute(ctx context.Context, keyspace, shard string, tabletType pbt.TabletType, query string, bindVars map[string]interface{}, transactionID int64) (<-chan *mproto.QueryResult, tabletconn.ErrFunc) {
var usedConn tabletconn.TabletConn
var erFunc tabletconn.ErrFunc
var results <-chan *mproto.QueryResult
err := dg.withRetry(ctx, keyspace, shard, tabletType, func(conn tabletconn.TabletConn) error {
var err error
results, erFunc, err = conn.StreamExecute2(ctx, query, bindVars, transactionID)
usedConn = conn
return err
}, transactionID, true)
if err != nil {
return results, func() error { return err }
}
inTransaction := (transactionID != 0)
return results, func() error {
return WrapError(erFunc(), keyspace, shard, tabletType, usedConn.EndPoint(), inTransaction)
}
}
// WrapError returns ShardConnError which preserves the original error code if possible,
// adds the connection context
// and adds a bit to determine whether the keyspace/shard needs to be
// re-resolved for a potential sharding event.
func (sdc *ShardConn) WrapError(in error, conn tabletconn.TabletConn, inTransaction bool) (wrapped error) {
if in == nil {
return nil
}
shardIdentifier := fmt.Sprintf("%s.%s.%s", sdc.keyspace, sdc.shard, sdc.tabletType)
if conn != nil {
shardIdentifier += fmt.Sprintf(", %+v", conn.EndPoint())
}
code := tabletconn.ERR_NORMAL
serverError, ok := in.(*tabletconn.ServerError)
if ok {
code = serverError.Code
}
shardConnErr := &ShardConnError{
Code: code,
ShardIdentifier: shardIdentifier,
InTransaction: inTransaction,
Err: in.Error(),
}
return shardConnErr
}