* Added ydb.WithOperationTimeout and ydb.WithOperationCancelAfter…

… context modifiers

CHANGELOG.md

@@ -1,8 +1,9 @@
-* Add go 1.18 to test matrix
+* Added `go1.18` to test matrix
+* Added `ydb.WithOperationTimeout` and `ydb.WithOperationCancelAfter` context modifiers
 
 ## v3.17.0
 * Removed redundant `trace.With{Table,Driver,Retry}` and `trace.Context{Table,Driver,Retry}` funcs 
-* Moved `gtrace` tool from `./cmt/gtrace` to `./internal/cmd/gtrace`
+* Moved `gtrace` tool from `./cmd/gtrace` to `./internal/cmd/gtrace`
 * Refactored `gtrace` tool for generate `Compose` options
 * Added panic recover on trace calls in `Compose` call step
 * Added `trace.With{Discovery,Driver,Coordination,Ratelimiter,Table,Scheme,Scripting}PanicCallback` options

context.go

@@ -0,0 +1,21 @@
+package ydb
+
+import (
+	"context"
+	"time"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/operation"
+)
+
+// WithOperationTimeout returns a copy of parent context in which YDB operation timeout
+// parameter is set to d. If parent context timeout is smaller than d, parent context is returned.
+func WithOperationTimeout(ctx context.Context, operationTimeout time.Duration) context.Context {
+	return operation.WithTimeout(ctx, operationTimeout)
+}
+
+// WithOperationCancelAfter returns a copy of parent context in which YDB operation
+// cancel after parameter is set to d. If parent context cancellation timeout is smaller
+// than d, parent context is returned.
+func WithOperationCancelAfter(ctx context.Context, operationCancelAfter time.Duration) context.Context {
+	return operation.WithCancelAfter(ctx, operationCancelAfter)
+}

internal/coordination/client.go

@@ -42,6 +42,7 @@ func (c *client) CreateNode(ctx context.Context, path string, config coordinatio
 				RateLimiterCountersMode:  config.RatelimiterCountersMode.To(),
 			},
 			OperationParams: operation.Params(
+				ctx,
 				c.config.OperationTimeout(),
 				c.config.OperationCancelAfter(),
 				operation.ModeSync,
@@ -65,6 +66,7 @@ func (c *client) AlterNode(ctx context.Context, path string, config coordination
 				RateLimiterCountersMode:  config.RatelimiterCountersMode.To(),
 			},
 			OperationParams: operation.Params(
+				ctx,
 				c.config.OperationTimeout(),
 				c.config.OperationCancelAfter(),
 				operation.ModeSync,
@@ -80,6 +82,7 @@ func (c *client) DropNode(ctx context.Context, path string) (err error) {
 		&Ydb_Coordination.DropNodeRequest{
 			Path: path,
 			OperationParams: operation.Params(
+				ctx,
 				c.config.OperationTimeout(),
 				c.config.OperationCancelAfter(),
 				operation.ModeSync,
@@ -107,6 +110,7 @@ func (c *client) DescribeNode(
 		&Ydb_Coordination.DescribeNodeRequest{
 			Path: path,
 			OperationParams: operation.Params(
+				ctx,
 				c.config.OperationTimeout(),
 				c.config.OperationCancelAfter(),
 				operation.ModeSync,

internal/operation/context.go

@@ -0,0 +1,54 @@
+package operation
+
+import (
+	"context"
+	"time"
+)
+
+type (
+	ctxOperationTimeoutKey     struct{}
+	ctxOperationCancelAfterKey struct{}
+)
+
+// WithTimeout returns a copy of parent context in which YDB operation timeout
+// parameter is set to d. If parent context timeout is smaller than d, parent context is returned.
+func WithTimeout(ctx context.Context, operationTimeout time.Duration) context.Context {
+	if d, ok := Timeout(ctx); ok && operationTimeout >= d {
+		// The current cancelation timeout is already smaller than the new one.
+		return ctx
+	}
+	return context.WithValue(ctx, ctxOperationTimeoutKey{}, operationTimeout)
+}
+
+// WithCancelAfter returns a copy of parent context in which YDB operation
+// cancel after parameter is set to d. If parent context cancellation timeout is smaller
+// than d, parent context is returned.
+func WithCancelAfter(ctx context.Context, operationCancelAfter time.Duration) context.Context {
+	if d, ok := CancelAfter(ctx); ok && operationCancelAfter >= d {
+		// The current cancelation timeout is already smaller than the new one.
+		return ctx
+	}
+	return context.WithValue(ctx, ctxOperationCancelAfterKey{}, operationCancelAfter)
+}
+
+// Timeout returns the timeout within given context after which
+// YDB should try to cancel operation and return result regardless of the cancelation.
+func Timeout(ctx context.Context) (d time.Duration, ok bool) {
+	d, ok = ctx.Value(ctxOperationTimeoutKey{}).(time.Duration)
+	return
+}
+
+// CancelAfter returns the timeout within given context after which
+// YDB should try to cancel operation and return result regardless of the cancellation.
+func CancelAfter(ctx context.Context) (d time.Duration, ok bool) {
+	d, ok = ctx.Value(ctxOperationCancelAfterKey{}).(time.Duration)
+	return
+}
+
+func untilDeadline(ctx context.Context) (time.Duration, bool) {
+	deadline, ok := ctx.Deadline()
+	if ok {
+		return time.Until(deadline), true
+	}
+	return 0, false
+}

internal/operation/params.go

@@ -1,16 +1,27 @@
 package operation
 
 import (
+	"context"
 	"time"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Operations"
 )
 
 func Params(
+	ctx context.Context,
 	timeout time.Duration,
 	cancelAfter time.Duration,
 	mode Mode,
 ) *Ydb_Operations.OperationParams {
+	if d, ok := Timeout(ctx); ok {
+		timeout = d
+	}
+	if d, ok := CancelAfter(ctx); ok {
+		cancelAfter = d
+	}
+	if d, ok := untilDeadline(ctx); mode == ModeSync && ok && d < timeout {
+		timeout = d
+	}
 	if timeout == 0 && cancelAfter == 0 && mode == 0 {
 		return nil
 	}