github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/operation/context.go (about)

     1  package operation
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  )
     7  
     8  type (
     9  	ctxOperationTimeoutKey     struct{}
    10  	ctxOperationCancelAfterKey struct{}
    11  )
    12  
    13  // WithTimeout returns a copy of parent context in which YDB operation timeout
    14  // parameter is set to d. If parent context timeout is smaller than d, parent context is returned.
    15  func WithTimeout(ctx context.Context, operationTimeout time.Duration) context.Context {
    16  	if d, ok := Timeout(ctx); ok && operationTimeout >= d {
    17  		// The current cancelation timeout is already smaller than the new one.
    18  		return ctx
    19  	}
    20  
    21  	return context.WithValue(ctx, ctxOperationTimeoutKey{}, operationTimeout)
    22  }
    23  
    24  // WithCancelAfter returns a copy of parent context in which YDB operation
    25  // cancel after parameter is set to d. If parent context cancellation timeout is smaller
    26  // than d, parent context is returned.
    27  func WithCancelAfter(ctx context.Context, operationCancelAfter time.Duration) context.Context {
    28  	if d, ok := CancelAfter(ctx); ok && operationCancelAfter >= d {
    29  		// The current cancelation timeout is already smaller than the new one.
    30  		return ctx
    31  	}
    32  
    33  	return context.WithValue(ctx, ctxOperationCancelAfterKey{}, operationCancelAfter)
    34  }
    35  
    36  // Timeout returns the timeout within given context after which
    37  // YDB should try to cancel operation and return result regardless of the cancelation.
    38  func Timeout(ctx context.Context) (d time.Duration, ok bool) {
    39  	d, ok = ctx.Value(ctxOperationTimeoutKey{}).(time.Duration)
    40  
    41  	return
    42  }
    43  
    44  // CancelAfter returns the timeout within given context after which
    45  // YDB should try to cancel operation and return result regardless of the cancellation.
    46  func CancelAfter(ctx context.Context) (d time.Duration, ok bool) {
    47  	d, ok = ctx.Value(ctxOperationCancelAfterKey{}).(time.Duration)
    48  
    49  	return
    50  }
    51  
    52  func untilDeadline(ctx context.Context) (time.Duration, bool) {
    53  	deadline, ok := ctx.Deadline()
    54  	if ok {
    55  		return time.Until(deadline), true
    56  	}
    57  
    58  	return 0, false
    59  }