Add my authorship

conn/inner_timeout.go

@@ -0,0 +1,51 @@
+package conn
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"google.golang.org/grpc"
+)
+
+type TimeoutError struct {
+	wrapped error
+}
+
+func (t *TimeoutError) Error() string {
+	return fmt.Sprintf("gosdk request timeout: %s", t.wrapped.Error())
+}
+
+func (t *TimeoutError) Unwrap() error {
+	return t.wrapped
+}
+
+// UnaryClientTimeoutInterceptor returns a new unary client interceptor that sets a timeout on the request context.
+func UnaryClientTimeoutInterceptor(timeout time.Duration) grpc.UnaryClientInterceptor {
+	return func(
+		ctx context.Context,
+		method string,
+		req, reply any,
+		cc *grpc.ClientConn,
+		invoker grpc.UnaryInvoker,
+		opts ...grpc.CallOption,
+	) error {
+		timedCtx, cancel := context.WithTimeout(ctx, timeout)
+		defer cancel()
+		err := invoker(timedCtx, method, req, reply, cc, opts...)
+		if err != nil {
+			if errors.Is(err, context.DeadlineExceeded) {
+				select {
+				case <-ctx.Done():
+					return err
+				default:
+					return &TimeoutError{wrapped: err}
+				}
+			}
+			// Otherwise, return the original error
+			return err
+		}
+		return nil
+	}
+}

operations/operation.go

@@ -15,6 +15,7 @@ import (
 	"google.golang.org/protobuf/types/known/anypb"
 	"google.golang.org/protobuf/types/known/emptypb"
 
+	"github.com/nebius/gosdk/conn"
 	common "github.com/nebius/gosdk/proto/nebius/common/v1"
 )
 
@@ -186,6 +187,10 @@ func (o *opWrapper) Wait(ctx context.Context, opts ...grpc.CallOption) (Operatio
 		case <-time.After(interval):
 			_, err := o.Poll(ctx, opts...)
 			if err != nil {
+				tErr := &conn.TimeoutError{}
+				if errors.As(err, &tErr) {
+					continue
+				}
 				return nil, err
 			}
 

sdk.go

@@ -9,7 +9,6 @@ import (
 	"time"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
-	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/timeout"
 	"google.golang.org/grpc"
 	grpc_creds "google.golang.org/grpc/credentials"
 
@@ -160,7 +159,7 @@ func New(ctx context.Context, opts ...Option) (*SDK, error) { //nolint:funlen
 	}
 
 	// apply timeout after retry and auth interceptors for each request
-	dialOpts = append(dialOpts, grpc.WithChainUnaryInterceptor(timeout.UnaryClientInterceptor(1*time.Minute)))
+	dialOpts = append(dialOpts, grpc.WithChainUnaryInterceptor(conn.UnaryClientTimeoutInterceptor(1*time.Minute)))
 
 	dialer := conn.NewDialer(logger, dialOpts...)
 	closes = append(closes, dialer.Close)