Golang callinfo.RPCWrapCallInfo函数代码示例

// ReloadSchema wraps RPCAgent.ReloadSchema
func (tm *TabletManager) ReloadSchema(ctx context.Context, args *rpc.Unused, reply *rpc.Unused) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrapLockAction(ctx, actionnode.TabletActionReloadSchema, args, reply, true, func() error {
		tm.agent.ReloadSchema(ctx)
		return nil
	})
}

// Backup wraps RPCAgent.Backup
func (tm *TabletManager) Backup(ctx context.Context, args *gorpcproto.BackupArgs, sendReply func(interface{}) error) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrapLockAction(ctx, actionnode.TabletActionBackup, args, nil, true, func() error {
		// create a logger, send the result back to the caller
		logger := logutil.NewChannelLogger(10)
		wg := sync.WaitGroup{}
		wg.Add(1)
		go func() {
			for e := range logger {
				// Note we don't interrupt the loop here, as
				// we still need to flush and finish the
				// command, even if the channel to the client
				// has been broken. We'll just keep trying
				// to send.
				sendReply(&e)
			}
			wg.Done()
		}()

		err := tm.agent.Backup(ctx, args.Concurrency, logger)
		close(logger)
		wg.Wait()
		return err
	})
}

// Ping wraps RPCAgent.Ping
func (tm *TabletManager) Ping(ctx context.Context, args, reply *string) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrap(ctx, actionnode.TabletActionPing, args, reply, func() error {
		*reply = tm.agent.Ping(ctx, *args)
		return nil
	})
}

// Sleep wraps RPCAgent.Sleep
func (tm *TabletManager) Sleep(ctx context.Context, args *time.Duration, reply *rpc.Unused) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrapLockAction(ctx, actionnode.TabletActionSleep, args, reply, true, func() error {
		tm.agent.Sleep(ctx, *args)
		return nil
	})
}

// ExecuteHook wraps RPCAgent.ExecuteHook
func (tm *TabletManager) ExecuteHook(ctx context.Context, args *hook.Hook, reply *hook.HookResult) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrapLockAction(ctx, actionnode.TabletActionExecuteHook, args, reply, true, func() error {
		*reply = *tm.agent.ExecuteHook(ctx, args)
		return nil
	})
}

// StreamExecute2 should not be used by anything other than tests.
// It will eventually replace Rollback, but it breaks compatibility with older clients.
// Once all clients are upgraded, it can be replaced.
func (sq *SqlQuery) StreamExecute2(ctx context.Context, req *proto.StreamExecuteRequest, sendReply func(reply interface{}) error) (err error) {
	defer sq.server.HandlePanic(&err)
	if req == nil || req.Query == nil {
		return nil
	}
	ctx = callerid.NewContext(ctx,
		callerid.GoRPCEffectiveCallerID(req.EffectiveCallerID),
		callerid.GoRPCImmediateCallerID(req.ImmediateCallerID),
	)
	tErr := sq.server.StreamExecute(callinfo.RPCWrapCallInfo(ctx), proto.TargetToProto3(req.Target), req.Query, func(reply *mproto.QueryResult) error {
		return sendReply(reply)
	})
	if tErr == nil {
		return nil
	}
	if *tabletserver.RPCErrorOnlyInReply {
		// If there was an app error, send a QueryResult back with it.
		qr := new(mproto.QueryResult)
		tabletserver.AddTabletErrorToQueryResult(tErr, qr)
		// Sending back errors this way is not backwards compatible. If a (new) server sends an additional
		// QueryResult with an error, and the (old) client doesn't know how to read it, it will cause
		// problems where the client will get out of sync with the number of QueryResults sent.
		// That's why this the error is only sent this way when the --rpc_errors_only_in_reply flag is set
		// (signalling that all clients are able to handle new-style errors).
		return sendReply(qr)
	}
	return tErr
}

// RunHealthCheck wraps RPCAgent.RunHealthCheck
func (tm *TabletManager) RunHealthCheck(ctx context.Context, args *pb.TabletType, reply *rpc.Unused) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrap(ctx, actionnode.TabletActionRunHealthCheck, args, reply, func() error {
		tm.agent.RunHealthCheck(ctx, *args)
		return nil
	})
}

// HealthStream registers an agent health stream
func (tm *TabletManager) HealthStream(ctx context.Context, args *rpc.Unused, sendReply func(interface{}) error) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrap(ctx, actionnode.TabletActionHealthStream, args, nil, func() error {
		c := make(chan *actionnode.HealthStreamReply, 10)
		wg := sync.WaitGroup{}
		wg.Add(1)
		go func() {
			defer wg.Done()
			for hsr := range c {
				// we send until the client disconnects
				if err := sendReply(hsr); err != nil {
					return
				}
			}
		}()

		id, err := tm.agent.RegisterHealthStream(c)
		if err != nil {
			close(c)
			wg.Wait()
			return err
		}
		wg.Wait()
		return tm.agent.UnregisterHealthStream(id)
	})
}

// GetSlaves wraps RPCAgent.GetSlaves
func (tm *TabletManager) GetSlaves(ctx context.Context, args *rpc.Unused, reply *gorpcproto.GetSlavesReply) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrap(ctx, actionnode.TabletActionGetSlaves, args, reply, func() error {
		var err error
		reply.Addrs, err = tm.agent.GetSlaves(ctx)
		return err
	})
}

// ExecuteBatch is exposing tabletserver.SqlQuery.ExecuteBatch
func (sq *SqlQuery) ExecuteBatch(ctx context.Context, queryList *proto.QueryList, reply *proto.QueryResultList) (err error) {
	defer sq.server.HandlePanic(&err)
	tErr := sq.server.ExecuteBatch(callinfo.RPCWrapCallInfo(ctx), nil, queryList, reply)
	tabletserver.AddTabletErrorToQueryResultList(tErr, reply)
	if *tabletserver.RPCErrorOnlyInReply {
		return nil
	}
	return tErr
}

// SplitQuery is exposing tabletserver.SqlQuery.SplitQuery
func (sq *SqlQuery) SplitQuery(ctx context.Context, req *proto.SplitQueryRequest, reply *proto.SplitQueryResult) (err error) {
	defer sq.server.HandlePanic(&err)
	tErr := sq.server.SplitQuery(callinfo.RPCWrapCallInfo(ctx), req, reply)
	tabletserver.AddTabletErrorToSplitQueryResult(tErr, reply)
	if *tabletserver.RPCErrorOnlyInReply {
		return nil
	}
	return tErr
}

// Begin is exposing tabletserver.SqlQuery.Begin
func (sq *SqlQuery) Begin(ctx context.Context, session *proto.Session, txInfo *proto.TransactionInfo) (err error) {
	defer sq.server.HandlePanic(&err)
	tErr := sq.server.Begin(callinfo.RPCWrapCallInfo(ctx), nil, session, txInfo)
	tabletserver.AddTabletErrorToTransactionInfo(tErr, txInfo)
	if *tabletserver.RPCErrorOnlyInReply {
		return nil
	}
	return tErr
}

// PromoteSlaveWhenCaughtUp wraps RPCAgent.PromoteSlaveWhenCaughtUp
func (tm *TabletManager) PromoteSlaveWhenCaughtUp(ctx context.Context, args *myproto.ReplicationPosition, reply *myproto.ReplicationPosition) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrapLockAction(ctx, actionnode.TabletActionPromoteSlaveWhenCaughtUp, args, reply, true, func() error {
		rp, err := tm.agent.PromoteSlaveWhenCaughtUp(ctx, *args)
		if err == nil {
			*reply = rp
		}
		return err
	})
}

// ApplySchema wraps RPCAgent.ApplySchema
func (tm *TabletManager) ApplySchema(ctx context.Context, args *myproto.SchemaChange, reply *myproto.SchemaChangeResult) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrapLockAction(ctx, actionnode.TabletActionApplySchema, args, reply, true, func() error {
		scr, err := tm.agent.ApplySchema(ctx, args)
		if err == nil {
			*reply = *scr
		}
		return err
	})
}

// InitMaster wraps RPCAgent.InitMaster
func (tm *TabletManager) InitMaster(ctx context.Context, args *rpc.Unused, reply *myproto.ReplicationPosition) error {
	ctx = callinfo.RPCWrapCallInfo(ctx)
	return tm.agent.RPCWrapLockAction(ctx, actionnode.TabletActionInitMaster, args, reply, true, func() error {
		rp, err := tm.agent.InitMaster(ctx)
		if err == nil {
			*reply = rp
		}
		return err
	})
}