feat: add context in s3 Get functions

Author: bpeng

aws/s3/s3.go

@@ -128,22 +128,32 @@ func (s *S3) Client() *s3.Client {
 // Get gets the object referred to by key and version from bucket and writes it into b.
 // Version can be empty.
 func (s *S3) Get(bucket, key, version string, b *bytes.Buffer) error {
+	return s.GetWithContext(context.Background(), bucket, key, version, b)
+}
+
+// Get gets the object referred to by key and version from bucket and writes it into b.
+// with the provided context.
+// Version can be empty.
+func (s *S3) GetWithContext(ctx context.Context, bucket, key, version string, w io.Writer) error {
 	input := s3.GetObjectInput{
 		Key:    aws.String(key),
 		Bucket: aws.String(bucket),
 	}
 	if version != "" {
 		input.VersionId = aws.String(version)
 	}
-	result, err := s.client.GetObject(context.TODO(), &input)
+	result, err := s.client.GetObject(ctx, &input)
 	if err != nil {
 		return err
 	}
 	defer result.Body.Close()
 
-	_, err = b.ReadFrom(result.Body)
-
+	_, err = io.Copy(w, result.Body)
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
 	return err
+
 }
 
 // GetByteRange gets the specified byte range of an object referred to by key and version

aws/s3/s3_concurrent.go

@@ -158,34 +158,91 @@ func newConcurrencyManager(maxWorkers, maxWorkersPerRequest, maxBytes int) *Conc
 // Version can be empty, but must be the same for all objects.
 func (s *S3Concurrent) GetAllConcurrently(bucket, version string, objects []types.Object) chan HydratedFile {
 
+	return s.GetAllConcurrentlyWithContext(context.Background(), bucket, version, objects)
+}
+
+// GetAllConcurrently gets the objects with provided context, from specified bucket and writes the resulting HydratedFiles
+// to the returned output channel. The closure of this channel is handled, however it's the caller's
+// responsibility to purge the channel, and handle any errors present in the HydratedFiles.
+// If the ConcurrencyManager is not initialised before calling GetAllConcurrently, an output channel
+// containing a single HydratedFile with an error is returned.
+// Version can be empty, but must be the same for all objects.
+func (s *S3Concurrent) GetAllConcurrentlyWithContext(
+	ctx context.Context,
+	bucket, version string,
+	objects []types.Object,
+) chan HydratedFile {
+
+	output := make(chan HydratedFile, 1)
+
+	// Early cancel check
+	select {
+	case <-ctx.Done():
+		output <- HydratedFile{Error: ctx.Err()}
+		close(output)
+		return output
+	default:
+	}
+
 	if s.manager == nil {
-		output := make(chan HydratedFile, 1)
-		output <- HydratedFile{Error: errors.New("error getting files from S3, Concurrency Manager not initialised")}
+		output <- HydratedFile{
+			Error: errors.New("error getting files from S3, Concurrency Manager not initialised"),
+		}
 		close(output)
 		return output
 	}
 
 	if s.manager.memoryTotalSize < s.manager.calculateRequiredMemoryFor(objects) {
-		output := make(chan HydratedFile, 1)
-		output <- HydratedFile{Error: fmt.Errorf("error: bytes requested greater than max allowed by server (%v)", s.manager.memoryTotalSize)}
+		output <- HydratedFile{
+			Error: fmt.Errorf(
+				"error: bytes requested greater than max allowed by server (%v)",
+				s.manager.memoryTotalSize,
+			),
+		}
 		close(output)
 		return output
 	}
-	// Secure memory for all objects upfront.
-	s.manager.secureMemory(objects) // 0.
+
+	// Secure memory for all objects upfront
+	s.manager.secureMemory(objects)
+
+	// IMPORTANT: ensure memory is released if context cancels before processing finishes
+	go func() {
+		<-ctx.Done()
+		// Best-effort cleanup: release all secured memory
+		for _, o := range objects {
+			s.manager.releaseMemory(aws.ToInt64(o.Size))
+		}
+	}()
 
 	processFunc := func(input types.Object) HydratedFile {
+		// Respect cancellation before starting work
+		select {
+		case <-ctx.Done():
+			return HydratedFile{Error: ctx.Err()}
+		default:
+		}
+
 		buf := bytes.NewBuffer(make([]byte, 0, int(*input.Size)))
 		key := aws.ToString(input.Key)
-		err := s.Get(bucket, key, version, buf)
+
+		// Prefer context-aware S3 call if available
+		err := s.GetWithContext(ctx, bucket, key, version, buf)
+
+		// If context was cancelled during S3 read, surface that
+		if ctx.Err() != nil {
+			return HydratedFile{Error: ctx.Err()}
+		}
 
 		return HydratedFile{
 			Key:   key,
 			Data:  buf.Bytes(),
 			Error: err,
 		}
 	}
-	return s.manager.Process(processFunc, objects)
+
+	// Process already accepts a context internally, so pass it through
+	return s.manager.ProcessWithContext(ctx, processFunc, objects)
 }
 
 // getWorker retrieves a number of workers from the manager's worker pool.
@@ -259,6 +316,33 @@ func (cm *ConcurrencyManager) Process(asyncProcessor FileProcessor, objects []ty
 	return workerGroup.returnOutput() // 2.
 }
 
+// Functions for providing a fan-out/fan-in operation with provided context. Workers are taken from the
+// worker pool and added to a WorkerGroup. All workers are returned to the pool once
+// the jobs have finished.
+func (cm *ConcurrencyManager) ProcessWithContext(
+	ctx context.Context,
+	asyncProcessor FileProcessor,
+	objects []types.Object,
+) chan HydratedFile {
+
+	workerGroup := cm.newWorkerGroup(ctx, asyncProcessor, cm.maxWorkersPerRequest)
+
+	go func() {
+		for _, obj := range objects {
+			select {
+			case <-ctx.Done():
+				workerGroup.stopWork()
+				return
+			default:
+				workerGroup.addWork(obj)
+			}
+		}
+		workerGroup.stopWork()
+	}()
+
+	return workerGroup.returnOutput()
+}
+
 // start begins a worker's process of making itself available for work, doing the work,
 // and repeat, until all work is done.
 func (w *worker) start(ctx context.Context, processor FileProcessor, roster chan *worker, wg *sync.WaitGroup) {