Golang tabletconn.GetDialer函数代码示例

// NewQueryResultReaderForTablet creates a new QueryResultReader for
// the provided tablet / sql query
func NewQueryResultReaderForTablet(ctx context.Context, ts topo.Server, tabletAlias topo.TabletAlias, sql string) (*QueryResultReader, error) {
	tablet, err := ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return nil, err
	}

	endPoint, err := tablet.EndPoint()
	if err != nil {
		return nil, err
	}

	conn, err := tabletconn.GetDialer()(ctx, *endPoint, tablet.Keyspace, tablet.Shard, *remoteActionsTimeout)
	if err != nil {
		return nil, err
	}

	sr, clientErrFn, err := conn.StreamExecute(ctx, sql, make(map[string]interface{}), 0)
	if err != nil {
		return nil, err
	}

	// read the columns, or grab the error
	cols, ok := <-sr
	if !ok {
		return nil, fmt.Errorf("Cannot read Fields for query '%v': %v", sql, clientErrFn())
	}

	return &QueryResultReader{
		Output:      sr,
		Fields:      cols.Fields,
		conn:        conn,
		clientErrFn: clientErrFn,
	}, nil
}

func commandVtTabletRollback(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() != 2 {
		return fmt.Errorf("the <tablet_alias> and <transaction_id> arguments are required for the VtTabletRollback command")
	}
	transactionID, err := strconv.ParseInt(subFlags.Arg(1), 10, 64)
	if err != nil {
		return err
	}
	tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
	if err != nil {
		return err
	}
	tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}

	conn, err := tabletconn.GetDialer()(tabletInfo.Tablet, *connectTimeout)
	if err != nil {
		return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
	}
	defer conn.Close(ctx)

	return conn.Rollback(ctx, &querypb.Target{
		Keyspace:   tabletInfo.Tablet.Keyspace,
		Shard:      tabletInfo.Tablet.Shard,
		TabletType: tabletInfo.Tablet.Type,
	}, transactionID)
}

// 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
}

// getTablet (re)sets the tablet which is used for the streaming query.
// If the method returns an error, the first return value specifies if it is
// okay to retry.
func (r *RestartableResultReader) getTablet() (bool, error) {
	if r.tablet != nil {
		// If there was a tablet before, return it to the tabletProvider.
		r.Close()

		r.tablet = nil
		r.conn = nil
		r.fields = nil
		r.output = nil
	}

	// Get a tablet from the tablet provider.
	tablet, err := r.tp.getTablet()
	if err != nil {
		return true /* retryable */, fmt.Errorf("failed get tablet for streaming query: %v", err)
	}

	// Get the dialer for it.
	conn, err := tabletconn.GetDialer()(tablet, *remoteActionsTimeout)
	if err != nil {
		return false /* retryable */, fmt.Errorf("failed to get dialer for tablet: %v", err)
	}
	r.tablet = tablet
	r.conn = conn
	return false /* retryable */, nil
}

// NewRestartableResultReader creates a new RestartableResultReader for
// the provided tablet and chunk.
// It will automatically create the necessary query to read all rows within
// the chunk.
// NOTE: We assume that the Columns field in "td" was ordered by a preceding
// call to reorderColumnsPrimaryKeyFirst().
func NewRestartableResultReader(ctx context.Context, logger logutil.Logger, ts topo.Server, tabletAlias *topodatapb.TabletAlias, td *tabletmanagerdatapb.TableDefinition, chunk chunk) (*RestartableResultReader, error) {
	shortCtx, cancel := context.WithTimeout(ctx, *remoteActionsTimeout)
	tablet, err := ts.GetTablet(shortCtx, tabletAlias)
	cancel()
	if err != nil {
		return nil, fmt.Errorf("tablet=%v table=%v chunk=%v: Failed to resolve tablet alias: %v", topoproto.TabletAliasString(tabletAlias), td.Name, chunk, err)
	}

	conn, err := tabletconn.GetDialer()(tablet.Tablet, *remoteActionsTimeout)
	if err != nil {
		return nil, fmt.Errorf("tablet=%v table=%v chunk=%v: Failed to get dialer for tablet: %v", topoproto.TabletAliasString(tabletAlias), td.Name, chunk, err)
	}

	r := &RestartableResultReader{
		ctx:    ctx,
		logger: logger,
		tablet: tablet.Tablet,
		td:     td,
		chunk:  chunk,
		conn:   conn,
	}

	if err := r.startStream(); err != nil {
		return nil, err
	}
	logger.Infof("tablet=%v table=%v chunk=%v: Starting to stream rows using query '%v'.", topoproto.TabletAliasString(tabletAlias), td.Name, chunk, r.query)
	return r, nil
}

// TestErrorSuite runs all the tests that expect errors
func TestErrorSuite(t *testing.T, protocol string, endPoint *topodatapb.EndPoint, fake *FakeQueryService) {
	// make sure we use the right client
	*tabletconn.TabletProtocol = protocol

	// create a connection, using sessionId
	ctx := context.Background()
	conn, err := tabletconn.GetDialer()(ctx, endPoint, testTarget.Keyspace, testTarget.Shard, topodatapb.TabletType_UNKNOWN, 30*time.Second)
	if err != nil {
		t.Fatalf("dial failed: %v", err)
	}

	// fake should return an error, make sure errors are handled properly
	fake.hasError = true
	testBeginError(t, conn)
	testCommitError(t, conn)
	testRollbackError(t, conn)
	testExecuteError(t, conn)
	testStreamExecuteError(t, conn, fake)
	testExecuteBatchError(t, conn)
	testSplitQueryError(t, conn)

	testBegin2Error(t, conn)
	testCommit2Error(t, conn)
	testRollback2Error(t, conn)
	testExecute2Error(t, conn)
	testStreamExecute2Error(t, conn, fake)
	testExecuteBatch2Error(t, conn)
	fake.hasError = false

	conn.Close()
}

// addL2VTGateConn adds a backend l2vtgate for the provided keyspace / shard.
func (lg *l2VTGateGateway) addL2VTGateConn(addr, keyspace, shard string) error {
	lg.mu.Lock()
	defer lg.mu.Unlock()

	// extract keyrange if it's a range
	canonical, kr, err := topo.ValidateShardName(shard)
	if err != nil {
		return fmt.Errorf("error parsing shard name %v: %v", shard, err)
	}

	// check for duplicates
	for _, c := range lg.connMap[keyspace] {
		if c.shard == canonical {
			return fmt.Errorf("duplicate %v/%v entry", keyspace, shard)
		}
	}

	// Dial in the background
	conn, err := tabletconn.GetDialer()(&topodatapb.Tablet{
		Hostname: addr,
	}, 0)
	if err != nil {
		return err
	}

	lg.connMap[keyspace] = append(lg.connMap[keyspace], &l2VTGateConn{
		addr:     addr,
		keyspace: keyspace,
		shard:    canonical,
		keyRange: kr,
		conn:     conn,
	})
	return nil
}

// NewQueryResultReaderForTablet creates a new QueryResultReader for
// the provided tablet / sql query
func NewQueryResultReaderForTablet(ctx context.Context, ts topo.Server, tabletAlias *topodatapb.TabletAlias, sql string) (*QueryResultReader, error) {
	shortCtx, cancel := context.WithTimeout(ctx, *remoteActionsTimeout)
	tablet, err := ts.GetTablet(shortCtx, tabletAlias)
	cancel()
	if err != nil {
		return nil, err
	}

	conn, err := tabletconn.GetDialer()(tablet.Tablet, *remoteActionsTimeout)
	if err != nil {
		return nil, err
	}

	stream, err := conn.StreamExecute(ctx, &querypb.Target{
		Keyspace:   tablet.Tablet.Keyspace,
		Shard:      tablet.Tablet.Shard,
		TabletType: tablet.Tablet.Type,
	}, sql, make(map[string]interface{}), nil)
	if err != nil {
		return nil, err
	}

	// read the columns, or grab the error
	cols, err := stream.Recv()
	if err != nil {
		return nil, fmt.Errorf("Cannot read Fields for query '%v': %v", sql, err)
	}

	return &QueryResultReader{
		output: stream,
		fields: cols.Fields,
		conn:   conn,
	}, nil
}

func commandVtTabletBegin(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() != 1 {
		return fmt.Errorf("the <tablet_alias> argument is required for the VtTabletBegin command")
	}
	tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
	if err != nil {
		return err
	}
	tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}

	conn, err := tabletconn.GetDialer()(ctx, tabletInfo.Tablet, *connectTimeout)
	if err != nil {
		return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
	}
	defer conn.Close()

	transactionID, err := conn.Begin(ctx)
	if err != nil {
		return fmt.Errorf("Begin failed: %v", err)
	}
	result := map[string]int64{
		"transaction_id": transactionID,
	}
	return printJSON(wr.Logger(), result)
}

func (th *tabletHealth) stream(ctx context.Context, ts topo.Server, tabletAlias *topodatapb.TabletAlias) (err error) {
	defer func() {
		th.mu.Lock()
		th.err = err
		th.mu.Unlock()
		close(th.done)
	}()

	ti, err := ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}
	ep, err := topo.TabletEndPoint(ti.Tablet)
	if err != nil {
		return err
	}

	// Pass in a tablet type that is not UNKNOWN, so we don't ask
	// for sessionId.
	conn, err := tabletconn.GetDialer()(ctx, ep, "", "", topodatapb.TabletType_MASTER, 30*time.Second)
	if err != nil {
		return err
	}
	defer conn.Close()

	stream, err := conn.StreamHealth(ctx)
	if err != nil {
		return err
	}

	first := true
	for time.Since(th.lastAccessed()) < *tabletHealthKeepAlive {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		result, err := stream.Recv()
		if err != nil {
			return err
		}

		th.mu.Lock()
		th.result = result
		th.mu.Unlock()

		if first {
			// We got the first result, so we're ready to be accessed.
			close(th.ready)
			first = false
		}
	}

	return nil
}

// getNewConn creates a new tablet connection with a separate per conn timeout.
// It limits the overall timeout to connTimeoutTotal by checking elapsed time after each blocking call.
func (sdc *ShardConn) getNewConn(ctx context.Context) (conn tabletconn.TabletConn, endPoint *topodatapb.EndPoint, isTimeout bool, err error) {
	startTime := time.Now()

	endPoints, err := sdc.balancer.Get()
	if err != nil {
		// Error when getting endpoint
		return nil, nil, false, err
	}
	if len(endPoints) == 0 {
		// No valid endpoint
		return nil, nil, false, vterrors.FromError(
			vtrpcpb.ErrorCode_INTERNAL_ERROR,
			fmt.Errorf("no valid endpoint"),
		)
	}
	if time.Now().Sub(startTime) >= sdc.connTimeoutTotal {
		return nil, nil, true, vterrors.FromError(
			vtrpcpb.ErrorCode_DEADLINE_EXCEEDED,
			fmt.Errorf("timeout when getting endpoints"),
		)
	}

	// Iterate through all endpoints to create a connection
	perConnTimeout := sdc.getConnTimeoutPerConn(len(endPoints))
	allErrors := new(concurrency.AllErrorRecorder)
	for _, endPoint := range endPoints {
		perConnStartTime := time.Now()
		conn, err = tabletconn.GetDialer()(ctx, endPoint, sdc.keyspace, sdc.shard, topodatapb.TabletType_UNKNOWN, perConnTimeout)
		if err == nil {
			sdc.connectTimings.Record([]string{sdc.keyspace, sdc.shard, strings.ToLower(sdc.tabletType.String())}, perConnStartTime)
			sdc.mu.Lock()
			defer sdc.mu.Unlock()
			sdc.conn = conn
			return conn, endPoint, false, nil
		}
		// Markdown the endpoint if it failed to connect
		sdc.balancer.MarkDown(endPoint.Uid, err.Error())
		vtErr := vterrors.NewVitessError(
			// TODO(aaijazi): what about OperationalErrors here?
			vterrors.RecoverVtErrorCode(err), err,
			"%v %+v", err, endPoint,
		)
		allErrors.RecordError(vtErr)
		if time.Now().Sub(startTime) >= sdc.connTimeoutTotal {
			err = vterrors.FromError(
				vtrpcpb.ErrorCode_DEADLINE_EXCEEDED,
				fmt.Errorf("timeout when connecting to %+v", endPoint),
			)
			allErrors.RecordError(err)
			return nil, nil, true, allErrors.AggrError(AggregateVtGateErrors)
		}
	}
	return nil, nil, false, allErrors.Error()
}

func (th *tabletHealth) stream(ctx context.Context, ts topo.Server, tabletAlias topo.TabletAlias) (err error) {
	defer func() {
		th.mu.Lock()
		th.err = err
		th.mu.Unlock()
		close(th.done)
	}()

	ti, err := ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}
	ep, err := ti.EndPoint()
	if err != nil {
		return err
	}

	// pass in empty keyspace and shard to not ask for sessionId
	conn, err := tabletconn.GetDialer()(ctx, *ep, "", "", 30*time.Second)
	if err != nil {
		return err
	}
	defer conn.Close()

	stream, errFunc, err := conn.StreamHealth(ctx)
	if err != nil {
		return err
	}

	first := true
	for time.Since(th.lastAccessed()) < *tabletHealthKeepAlive {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case result, ok := <-stream:
			if !ok {
				return errFunc()
			}

			th.mu.Lock()
			th.result = result
			th.mu.Unlock()

			if first {
				// We got the first result, so we're ready to be accessed.
				close(th.ready)
				first = false
			}
		}
	}

	return nil
}

func (th *tabletHealth) stream(ctx context.Context, ts topo.Server, tabletAlias *topodatapb.TabletAlias) (err error) {
	defer func() {
		th.mu.Lock()
		th.err = err
		th.mu.Unlock()
		close(th.done)
	}()

	ti, err := ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}

	conn, err := tabletconn.GetDialer()(ti.Tablet, 30*time.Second)
	if err != nil {
		return err
	}
	defer conn.Close(ctx)

	stream, err := conn.StreamHealth(ctx)
	if err != nil {
		return err
	}

	first := true
	for time.Since(th.lastAccessed()) < *tabletHealthKeepAlive {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		result, err := stream.Recv()
		if err != nil {
			return err
		}

		th.mu.Lock()
		th.result = result
		th.mu.Unlock()

		if first {
			// We got the first result, so we're ready to be accessed.
			close(th.ready)
			first = false
		}
	}

	return nil
}

// connect creates connection to the endpoint and starts streaming.
func (hcc *healthCheckConn) connect(ctx context.Context, hc *HealthCheckImpl, endPoint *pbt.EndPoint) (<-chan *pbq.StreamHealthResponse, tabletconn.ErrFunc, error) {
	conn, err := tabletconn.GetDialer()(ctx, endPoint, "" /*keyspace*/, "" /*shard*/, pbt.TabletType_RDONLY, hc.connTimeout)
	if err != nil {
		return nil, nil, err
	}
	stream, errfunc, err := conn.StreamHealth(ctx)
	if err != nil {
		conn.Close()
		return nil, nil, err
	}
	hcc.mu.Lock()
	hcc.conn = conn
	hcc.mu.Unlock()
	return stream, errfunc, nil
}

func commandVtTabletExecute(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
	transactionID := subFlags.Int("transaction_id", 0, "transaction id to use, if inside a transaction.")
	bindVariables := newBindvars(subFlags)
	keyspace := subFlags.String("keyspace", "", "keyspace the tablet belongs to")
	shard := subFlags.String("shard", "", "shard the tablet belongs to")
	tabletType := subFlags.String("tablet_type", "unknown", "tablet type we expect from the tablet (use unknown to use sessionId)")
	connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
	json := subFlags.Bool("json", false, "Output JSON instead of human-readable table")

	if err := subFlags.Parse(args); err != nil {
		return err
	}
	if subFlags.NArg() != 2 {
		return fmt.Errorf("the <tablet_alias> and <sql> arguments are required for the VtTabletExecute command")
	}
	tt, err := topoproto.ParseTabletType(*tabletType)
	if err != nil {
		return err
	}
	tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
	if err != nil {
		return err
	}
	tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}
	ep, err := topo.TabletEndPoint(tabletInfo.Tablet)
	if err != nil {
		return fmt.Errorf("cannot get EndPoint from tablet record: %v", err)
	}

	conn, err := tabletconn.GetDialer()(ctx, ep, *keyspace, *shard, tt, *connectTimeout)
	if err != nil {
		return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
	}
	defer conn.Close()

	qr, err := conn.Execute(ctx, subFlags.Arg(1), *bindVariables, int64(*transactionID))
	if err != nil {
		return fmt.Errorf("Execute failed: %v", err)
	}
	if *json {
		return printJSON(wr.Logger(), qr)
	}
	printQueryResult(loggerWriter{wr.Logger()}, qr)
	return nil
}