feat: add timeout to individual requests in batch call

batchutils/batch_delete.go

@@ -3,6 +3,7 @@ package batchutils
 import (
 	"context"
 	"sync"
+	"time"
 
 	"github.com/momentohq/client-sdk-go/momento"
 )
@@ -15,16 +16,23 @@ func deleteWorker(
 	cacheName string,
 	keyChan chan momento.Key,
 	errChan chan *errKeyVal,
+	timeout time.Duration,
 ) {
 	for {
 		myKey := <-keyChan
 		if myKey == nil {
 			return
 		}
-		_, err := client.Delete(ctx, &momento.DeleteRequest{
+
+		timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
+
+		_, err := client.Delete(timeoutCtx, &momento.DeleteRequest{
 			CacheName: cacheName,
 			Key:       myKey,
 		})
+
+		cancel()
+
 		if err != nil {
 			errChan <- &errKeyVal{
 				key:   myKey,
@@ -41,6 +49,8 @@ type BatchDeleteRequest struct {
 	CacheName            string
 	Keys                 []momento.Key
 	MaxConcurrentDeletes int
+	// timeout for individual requests, defaults to 10 seconds
+	RequestTimeout *time.Duration
 }
 
 // BatchDeleteError contains a map associating failing cache keys with their specific errors.
@@ -81,7 +91,7 @@ func BatchDelete(ctx context.Context, props *BatchDeleteRequest) *BatchDeleteErr
 
 		go func() {
 			defer wg.Done()
-			deleteWorker(ctx, props.Client, props.CacheName, keyChan, errChan)
+			deleteWorker(ctx, props.Client, props.CacheName, keyChan, errChan, getRequestTimeout(props.RequestTimeout))
 		}()
 	}
 

batchutils/batch_get.go

@@ -3,6 +3,7 @@ package batchutils
 import (
 	"context"
 	"sync"
+	"time"
 
 	"github.com/momentohq/client-sdk-go/momento"
 	"github.com/momentohq/client-sdk-go/responses"
@@ -21,16 +22,23 @@ func getWorker(
 	cacheName string,
 	keyChan chan momento.Key,
 	resultChan chan *getResultOrError,
+	timeout time.Duration,
 ) {
 	for {
 		myKey := <-keyChan
 		if myKey == nil {
 			return
 		}
-		getResponse, err := client.Get(ctx, &momento.GetRequest{
+
+		timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
+
+		getResponse, err := client.Get(timeoutCtx, &momento.GetRequest{
 			CacheName: cacheName,
 			Key:       myKey,
 		})
+
+		cancel()
+
 		if err != nil {
 			resultChan <- &getResultOrError{err: &errKeyVal{
 				key:   myKey,
@@ -50,6 +58,8 @@ type BatchGetRequest struct {
 	CacheName         string
 	Keys              []momento.Key
 	MaxConcurrentGets int
+	// timeout for individual requests, defaults to 10 seconds
+	RequestTimeout *time.Duration
 }
 
 // BatchGetError contains a map associating failing cache keys with their specific errors.
@@ -104,7 +114,7 @@ func BatchGet(ctx context.Context, props *BatchGetRequest) (*BatchGetResponse, *
 
 		go func() {
 			defer wg.Done()
-			getWorker(ctx, props.Client, props.CacheName, keyChan, resultChan)
+			getWorker(ctx, props.Client, props.CacheName, keyChan, resultChan, getRequestTimeout(props.RequestTimeout))
 		}()
 	}
 

batchutils/batch_operations.go

@@ -2,10 +2,22 @@ package batchutils
 
 import (
 	"context"
+	"time"
 
 	"github.com/momentohq/client-sdk-go/momento"
 )
 
+const defaultRequestTimeout = 10 * time.Second
+
+func getRequestTimeout(propsTimeout *time.Duration) (requestTimeout time.Duration) {
+	if propsTimeout == nil {
+		requestTimeout = defaultRequestTimeout
+	} else {
+		requestTimeout = *propsTimeout
+	}
+	return
+}
+
 func keyDistributor(ctx context.Context, keys []momento.Key, keyChan chan momento.Key) {
 	for _, k := range keys {
 		keyChan <- k

batchutils/batch_operations_test.go

@@ -103,6 +103,21 @@ var _ = Describe("Batch operations", func() {
 			})
 			Expect(errors).To(BeNil())
 		})
+
+		It("super small request timeout test", func() {
+			timeout := 1 * time.Millisecond
+			errors := batchutils.BatchDelete(ctx, &batchutils.BatchDeleteRequest{
+				Client:         client,
+				CacheName:      cacheName,
+				Keys:           keys[5:21],
+				RequestTimeout: &timeout,
+			})
+
+			Expect(len(errors.Errors())).To(Equal(16))
+			for _, err := range errors.Errors() {
+				Expect(err.Error()).To(ContainSubstring("TimeoutError: context deadline exceeded"))
+			}
+		})
 	})
 
 	Describe("batch get", func() {
@@ -158,5 +173,21 @@ var _ = Describe("Batch operations", func() {
 				}
 			}
 		})
+
+		It("super small request timeout test", func() {
+			timeout := 1 * time.Millisecond
+			getBatch, errors := batchutils.BatchGet(ctx, &batchutils.BatchGetRequest{
+				Client:         client,
+				CacheName:      cacheName,
+				Keys:           keys[5:21],
+				RequestTimeout: &timeout,
+			})
+
+			Expect(getBatch).To(BeNil())
+			Expect(len(errors.Errors())).To(Equal(16))
+			for _, err := range errors.Errors() {
+				Expect(err.Error()).To(ContainSubstring("TimeoutError: context deadline exceeded"))
+			}
+		})
 	})
 })

batchutils/batch_set.go

@@ -22,19 +22,25 @@ func setWorker(
 	cacheName string,
 	itemChan chan BatchSetItem,
 	resultChan chan *setResultOrError,
+	timeout time.Duration,
 ) {
 	for {
 		item := <-itemChan
 		if item.Key == nil {
 			return
 		}
-		setResponse, err := client.Set(ctx, &momento.SetRequest{
+
+		timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
+
+		setResponse, err := client.Set(timeoutCtx, &momento.SetRequest{
 			CacheName: cacheName,
 			Key:       item.Key,
 			Value:     item.Value,
 			Ttl:       item.Ttl,
 		})
 
+		cancel()
+
 		if err != nil {
 			resultChan <- &setResultOrError{err: &errKeyVal{
 				key:   item.Key,
@@ -54,6 +60,8 @@ type BatchSetRequest struct {
 	CacheName         string
 	Items             []BatchSetItem
 	MaxConcurrentSets int
+	// timeout for individual requests, defaults to 10 seconds
+	RequestTimeout *time.Duration
 }
 
 type BatchSetItem struct {
@@ -130,7 +138,7 @@ func BatchSet(ctx context.Context, props *BatchSetRequest) (*BatchSetResponse, *
 
 		go func() {
 			defer wg.Done()
-			setWorker(ctx, props.Client, props.CacheName, itemChan, resultChan)
+			setWorker(ctx, props.Client, props.CacheName, itemChan, resultChan, getRequestTimeout(props.RequestTimeout))
 		}()
 	}
 

batchutils/batch_set_test.go

@@ -252,5 +252,35 @@ var _ = Describe("Batch set operations", func() {
 				Expect(resp).To(BeAssignableToTypeOf(&responses.GetMiss{}))
 			}
 		})
+
+		It("super small request timeout test", func() {
+			var items []batchutils.BatchSetItem
+
+			for i := 0; i < 10; i++ {
+				key := String(fmt.Sprintf("MSETk%d", i))
+				item := batchutils.BatchSetItem{
+					Key:   key,
+					Value: String(fmt.Sprintf("MSETv%d", i)),
+					Ttl:   1 * time.Second,
+				}
+				items = append(items, item)
+			}
+
+			timeout := 1 * time.Millisecond
+			setBatchResp, setErrors := batchutils.BatchSet(ctx, &batchutils.BatchSetRequest{
+				Client:         client,
+				CacheName:      cacheName,
+				Items:          items,
+				RequestTimeout: &timeout,
+			})
+
+			Expect(setBatchResp).To(BeNil())
+			Expect(len(setErrors.Errors())).To(Equal(len(items)))
+			for _, err := range setErrors.Errors() {
+				Expect(err.Error()).To(ContainSubstring("TimeoutError: context deadline exceeded"))
+			}
+
+		})
+
 	})
 })