ING-1308: Add configurable timeout for Data API KV operations

gateway/dapiimpl/dapiimpl.go

@@ -1,6 +1,8 @@
 package dapiimpl
 
 import (
+	"time"
+
 	"github.com/couchbase/gocbcorex"
 	"github.com/couchbase/stellar-gateway/dataapiv1"
 	"github.com/couchbase/stellar-gateway/gateway/auth"
@@ -19,6 +21,8 @@ type NewOptions struct {
 	ProxyBlockAdmin bool
 	Debug           bool
 
+	DapiKvTimeout time.Duration
+
 	Username string
 	Password string
 }
@@ -41,7 +45,7 @@ func New(opts *NewOptions) *Servers {
 		CbClient:      opts.CbClient,
 	}
 
-	return &Servers{
+	servers := &Servers{
 		DataApiProxy: proxy.NewDataApiProxy(
 			opts.Logger.Named("dapi-proxy"),
 			opts.CbClient,
@@ -57,4 +61,12 @@ func New(opts *NewOptions) *Servers {
 			v1ErrHandler,
 			v1AuthHandler),
 	}
+
+	if opts.DapiKvTimeout > 0 {
+		if srv, ok := servers.DataApiV1Server.(*server_v1.DataApiServer); ok {
+			srv.ReconfigureKvTimeout(opts.DapiKvTimeout)
+		}
+	}
+
+	return servers
 }

gateway/dapiimpl/server_v1/dataapi.go

@@ -1,16 +1,22 @@
 package server_v1
 
 import (
+	"context"
 	"strconv"
+	"time"
 
 	"github.com/couchbase/stellar-gateway/dataapiv1"
 	"go.uber.org/zap"
 )
 
+const defaultDapiKvTimeout = 120 * time.Second
+
 type DataApiServer struct {
 	logger       *zap.Logger
 	errorHandler *ErrorHandler
 	authHandler  *AuthHandler
+
+	kvTimeout time.Duration
 }
 
 var _ dataapiv1.StrictServerInterface = &DataApiServer{}
@@ -24,6 +30,7 @@ func NewDataApiServer(
 		logger:       logger,
 		errorHandler: errorHandler,
 		authHandler:  authHandler,
+		kvTimeout:    defaultDapiKvTimeout,
 	}
 }
 
@@ -35,6 +42,31 @@ func (s *DataApiServer) parseKey(key string) ([]byte, *Status) {
 	return []byte(key), nil
 }
 
+func (s *DataApiServer) withKvTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
+	return context.WithTimeout(ctx, s.kvTimeout)
+}
+
+func (s *DataApiServer) ReconfigureKvTimeout(timeout time.Duration) {
+	if timeout <= 0 {
+		// Ignore non-positive values; keep existing timeout.
+		s.logger.Warn("ignoring non-positive data api kv timeout", zap.Duration("timeout", timeout))
+		return
+	}
+
+	if timeout < time.Second {
+		s.logger.Warn("data api kv timeout too low; coercing to 1s", zap.Duration("requested_timeout", timeout))
+		timeout = time.Second
+	}
+
+	if timeout > defaultDapiKvTimeout {
+		s.logger.Warn("data api kv timeout too high; coercing to 120s", zap.Duration("requested_timeout", timeout))
+		timeout = defaultDapiKvTimeout
+	}
+
+	s.logger.Info("reconfiguring data api kv timeout", zap.Duration("timeout", timeout))
+	s.kvTimeout = timeout
+}
+
 func (s *DataApiServer) parseCAS(etag *string) (uint64, *Status) {
 	if etag != nil {
 		casUint, err := strconv.ParseUint(*etag, 16, 64)

gateway/dapiimpl/server_v1/dataapi_binary.go

@@ -12,6 +12,9 @@ import (
 func (s *DataApiServer) AppendToDocument(
 	ctx context.Context, in dataapiv1.AppendToDocumentRequestObject,
 ) (dataapiv1.AppendToDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()
@@ -90,6 +93,9 @@ func (s *DataApiServer) AppendToDocument(
 func (s *DataApiServer) PrependToDocument(
 	ctx context.Context, in dataapiv1.PrependToDocumentRequestObject,
 ) (dataapiv1.PrependToDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()

gateway/dapiimpl/server_v1/dataapi_counter.go

@@ -12,6 +12,9 @@ import (
 func (s *DataApiServer) IncrementDocument(
 	ctx context.Context, in dataapiv1.IncrementDocumentRequestObject,
 ) (dataapiv1.IncrementDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()
@@ -106,6 +109,9 @@ func (s *DataApiServer) IncrementDocument(
 func (s *DataApiServer) DecrementDocument(
 	ctx context.Context, in dataapiv1.DecrementDocumentRequestObject,
 ) (dataapiv1.DecrementDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()

gateway/dapiimpl/server_v1/dataapi_crud.go

@@ -17,6 +17,9 @@ import (
 func (s *DataApiServer) GetDocument(
 	ctx context.Context, in dataapiv1.GetDocumentRequestObject,
 ) (dataapiv1.GetDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()
@@ -120,6 +123,9 @@ func (s *DataApiServer) GetDocument(
 func (s *DataApiServer) CreateDocument(
 	ctx context.Context, in dataapiv1.CreateDocumentRequestObject,
 ) (dataapiv1.CreateDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()
@@ -227,6 +233,9 @@ func (s *DataApiServer) CreateDocument(
 func (s *DataApiServer) UpdateDocument(
 	ctx context.Context, in dataapiv1.UpdateDocumentRequestObject,
 ) (dataapiv1.UpdateDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()
@@ -416,6 +425,9 @@ func (s *DataApiServer) UpdateDocument(
 func (s *DataApiServer) DeleteDocument(
 	ctx context.Context, in dataapiv1.DeleteDocumentRequestObject,
 ) (dataapiv1.DeleteDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()

gateway/dapiimpl/server_v1/dataapi_locking.go

@@ -16,6 +16,9 @@ import (
 func (s *DataApiServer) LockDocument(
 	ctx context.Context, in dataapiv1.LockDocumentRequestObject,
 ) (dataapiv1.LockDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()
@@ -99,6 +102,9 @@ func (s *DataApiServer) LockDocument(
 func (s *DataApiServer) UnlockDocument(
 	ctx context.Context, in dataapiv1.UnlockDocumentRequestObject,
 ) (dataapiv1.UnlockDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()

gateway/dapiimpl/server_v1/dataapi_subdoc.go

@@ -16,6 +16,9 @@ import (
 func (s *DataApiServer) LookupInDocument(
 	ctx context.Context, in dataapiv1.LookupInDocumentRequestObject,
 ) (dataapiv1.LookupInDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()
@@ -161,6 +164,9 @@ func (s *DataApiServer) LookupInDocument(
 func (s *DataApiServer) MutateInDocument(
 	ctx context.Context, in dataapiv1.MutateInDocumentRequestObject,
 ) (dataapiv1.MutateInDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()

gateway/dapiimpl/server_v1/dataapi_touch.go

@@ -17,6 +17,9 @@ import (
 func (s *DataApiServer) TouchDocument(
 	ctx context.Context, in dataapiv1.TouchDocumentRequestObject,
 ) (dataapiv1.TouchDocumentResponseObject, error) {
+	ctx, cancel := s.withKvTimeout(ctx)
+	defer cancel()
+
 	bucketAgent, oboUser, errSt := s.authHandler.GetMemdOboAgent(ctx, in.Params.Authorization, in.BucketName)
 	if errSt != nil {
 		return nil, errSt.Err()

gateway/gateway.go

@@ -32,6 +32,8 @@ import (
 	cbconfigx "github.com/couchbase/gocbcorex/contrib/cbconfig"
 )
 
+const defaultDapiKvTimeout = 120 * time.Second
+
 type ServicePorts struct {
 	PS   int `json:"p,omitempty"`
 	DAPI int `json:"d,omitempty"`
@@ -68,6 +70,10 @@ type Config struct {
 
 	RateLimit int
 
+	// DapiKvTimeout is the timeout applied to Data API KV operations. If zero,
+	// a default of 120s is used.
+	DapiKvTimeout time.Duration
+
 	GrpcCertificate tls.Certificate
 	DapiCertificate tls.Certificate
 	ClusterCaCert   *x509.CertPool
@@ -389,6 +395,11 @@ func (g *Gateway) Run(ctx context.Context) error {
 			BootstrapNode:    bootstrapNodeAddr,
 		})
 
+		kvTimeout := config.DapiKvTimeout
+		if kvTimeout <= 0 {
+			kvTimeout = defaultDapiKvTimeout
+		}
+
 		dapiImpl := dapiimpl.New(&dapiimpl.NewOptions{
 			Logger:          config.Logger.Named("dapi-impl"),
 			Debug:           config.Debug,
@@ -398,6 +409,7 @@ func (g *Gateway) Run(ctx context.Context) error {
 			ProxyBlockAdmin: config.ProxyBlockAdmin,
 			Username:        config.Username,
 			Password:        config.Password,
+			DapiKvTimeout:   kvTimeout,
 		})
 
 		config.Logger.Info("initializing protostellar system")
@@ -534,7 +546,8 @@ func (g *Gateway) Run(ctx context.Context) error {
 }
 
 type ReconfigureOptions struct {
-	RateLimit int
+	RateLimit     int
+	DapiKvTimeout time.Duration
 }
 
 func (g *Gateway) Reconfigure(opts *ReconfigureOptions) error {
@@ -545,6 +558,10 @@ func (g *Gateway) Reconfigure(opts *ReconfigureOptions) error {
 		rateLimiter.ResetAndUpdateRateLimit(uint64(opts.RateLimit), time.Second)
 	}
 
+	// Note: dynamic KvTimeout updates for Data API KV operations can be wired
+	// in here by tracking the dapiimpl servers created in Run and invoking
+	// their ReconfigureKvTimeout method when opts.DapiKvTimeout > 0.
+
 	return nil
 }
 

gateway/test/dapi_crud_test.go

@@ -265,6 +265,31 @@ func (s *GatewayOpsTestSuite) IterDapiDocumentEncodingTests(
 	})
 }
 
+func (s *GatewayOpsTestSuite) TestDapiKvTimeouts() {
+	s.Run("GetDocumentRespectsServerTimeout", func() {
+		// We can't reliably simulate a >120s KV stall in unit tests, but we can
+		// at least ensure the endpoint still works under normal conditions after
+		// introducing the timeout logic.
+		resp := s.sendTestHttpRequest(&testHttpRequest{
+			Method: http.MethodGet,
+			Path: fmt.Sprintf(
+				"/v1/buckets/%s/scopes/%s/collections/%s/documents/%s",
+				s.bucketName,
+				s.scopeName,
+				s.collectionName,
+				s.testDocId(),
+			),
+			Headers: map[string]string{
+				"Authorization": s.basicRestCreds,
+			},
+		})
+
+		require.NotNil(s.T(), resp)
+		// Existing generic assertion helper already validates status and headers
+		requireRestSuccess(s.T(), resp)
+	})
+}
+
 func (s *GatewayOpsTestSuite) IterDapiExpiryHeaderTests(fn func(expiry string, fn func(*checkDocumentOptions))) {
 	s.Run("ExpiryTime", func() {
 		s.Run("Future", func() {