github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/xcontext/context_with_timeout.go (about)

     1  package xcontext
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
    10  )
    11  
    12  func WithTimeout(ctx context.Context, t time.Duration) (context.Context, context.CancelFunc) {
    13  	childCtx := &timeoutCtx{
    14  		parentCtx: ctx,
    15  		from:      stack.Record(1),
    16  	}
    17  	childCtx.ctx, childCtx.ctxCancel = context.WithTimeout(ctx, t)
    18  
    19  	return childCtx, childCtx.cancel
    20  }
    21  
    22  type timeoutCtx struct {
    23  	parentCtx context.Context //nolint:containedctx
    24  	ctx       context.Context //nolint:containedctx
    25  	ctxCancel context.CancelFunc
    26  	from      string
    27  
    28  	m   sync.Mutex
    29  	err error
    30  }
    31  
    32  func (ctx *timeoutCtx) Deadline() (deadline time.Time, ok bool) {
    33  	return ctx.ctx.Deadline()
    34  }
    35  
    36  func (ctx *timeoutCtx) Done() <-chan struct{} {
    37  	return ctx.ctx.Done()
    38  }
    39  
    40  func (ctx *timeoutCtx) Err() error {
    41  	ctx.m.Lock()
    42  	defer ctx.m.Unlock()
    43  
    44  	if ctx.err != nil {
    45  		return ctx.err
    46  	}
    47  
    48  	if errors.Is(ctx.ctx.Err(), context.DeadlineExceeded) && ctx.parentCtx.Err() == nil {
    49  		ctx.err = errFrom(context.DeadlineExceeded, ctx.from)
    50  
    51  		return ctx.err
    52  	}
    53  
    54  	if err := ctx.parentCtx.Err(); err != nil {
    55  		ctx.err = err
    56  
    57  		return ctx.err
    58  	}
    59  
    60  	return nil
    61  }
    62  
    63  func (ctx *timeoutCtx) Value(key interface{}) interface{} {
    64  	return ctx.ctx.Value(key)
    65  }
    66  
    67  func (ctx *timeoutCtx) cancel() {
    68  	ctx.m.Lock()
    69  	defer ctx.m.Unlock()
    70  
    71  	ctx.ctxCancel()
    72  
    73  	if ctx.err != nil {
    74  		return
    75  	}
    76  
    77  	if err := ctx.parentCtx.Err(); err != nil {
    78  		ctx.err = err
    79  
    80  		return
    81  	}
    82  	ctx.err = errAt(context.Canceled, 1)
    83  }