// RunLimitedAsyncTask runs function f in a goroutine, using the given
// channel as a semaphore to limit the number of tasks that are run
// concurrently to the channel's capacity. If wait is true, blocks
// until the semaphore is available in order to push back on callers
// that may be trying to create many tasks. If wait is false, returns
// immediately with an error if the semaphore is not
// available. Returns an error if the Stopper is quiescing, in which
// case the function is not executed.
func (s *Stopper) RunLimitedAsyncTask(
ctx context.Context, sem chan struct{}, wait bool, f func(context.Context),
) error {
file, line, _ := caller.Lookup(1)
key := taskKey{file, line}
// Wait for permission to run from the semaphore.
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
case <-s.ShouldQuiesce():
return errUnavailable
default:
if !wait {
return ErrThrottled
}
log.Infof(context.TODO(), "stopper throttling task from %s:%d due to semaphore", file, line)
// Retry the select without the default.
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
case <-s.ShouldQuiesce():
return errUnavailable
}
}
// Check for canceled context: it's possible to get the semaphore even
// if the context is canceled.
select {
case <-ctx.Done():
<-sem
return ctx.Err()
default:
}
if !s.runPrelude(key) {
<-sem
return errUnavailable
}
ctx, span := tracing.ForkCtxSpan(ctx, fmt.Sprintf("%s:%d", file, line))
go func() {
defer s.Recover()
defer s.runPostlude(key)
defer func() { <-sem }()
defer tracing.FinishSpan(span)
f(ctx)
}()
return nil
}
// SendNext invokes the specified RPC on the supplied client when the
// client is ready. On success, the reply is sent on the channel;
// otherwise an error is sent.
func (gt *grpcTransport) SendNext(ctx context.Context, done chan<- BatchCall) {
client := gt.orderedClients[gt.clientIndex]
gt.clientIndex++
gt.setPending(client.args.Replica, true)
// Fork the original context as this async send may outlast the
// caller's context.
ctx, sp := tracing.ForkCtxSpan(ctx, "grpcTransport SendNext")
go func() {
defer tracing.FinishSpan(sp)
gt.opts.metrics.SentCount.Inc(1)
reply, err := func() (*roachpb.BatchResponse, error) {
if enableLocalCalls {
if localServer := gt.rpcContext.GetLocalInternalServerForAddr(client.remoteAddr); localServer != nil {
// Clone the request. At the time of writing, Replica may mutate it
// during command execution which can lead to data races.
//
// TODO(tamird): we should clone all of client.args.Header, but the
// assertions in protoutil.Clone fire and there seems to be no
// reasonable workaround.
origTxn := client.args.Txn
if origTxn != nil {
clonedTxn := origTxn.Clone()
client.args.Txn = &clonedTxn
}
gt.opts.metrics.LocalSentCount.Inc(1)
log.VEvent(ctx, 2, "sending request to local server")
return localServer.Batch(ctx, &client.args)
}
}
log.VEventf(ctx, 2, "sending request to %s", client.remoteAddr)
reply, err := client.client.Batch(ctx, &client.args)
if reply != nil {
for i := range reply.Responses {
if err := reply.Responses[i].GetInner().Verify(client.args.Requests[i].GetInner()); err != nil {
log.Error(ctx, err)
}
}
}
return reply, err
}()
gt.setPending(client.args.Replica, false)
done <- BatchCall{Reply: reply, Err: err}
}()
}
// RunAsyncTask runs function f in a goroutine. It returns an error when the
// Stopper is quiescing, in which case the function is not executed.
func (s *Stopper) RunAsyncTask(ctx context.Context, f func(context.Context)) error {
file, line, _ := caller.Lookup(1)
key := taskKey{file, line}
if !s.runPrelude(key) {
return errUnavailable
}
ctx, span := tracing.ForkCtxSpan(ctx, fmt.Sprintf("%s:%d", file, line))
// Call f.
go func() {
defer s.Recover()
defer s.runPostlude(key)
defer tracing.FinishSpan(span)
f(ctx)
}()
return nil
}
