update with NewBatchFuture api

Author: shijiesheng

internal/batch/__debug_bin2377381568

@@ -0,0 +1,149 @@
+package batch
+
+import (
+	"fmt"
+	"reflect"
+
+	"go.uber.org/multierr"
+
+	"go.uber.org/cadence/internal"
+)
+
+type BatchFuture interface {
+	internal.Future
+	GetFutures() []internal.Future
+}
+
+type batchFutureImpl struct {
+	futures   []internal.Future
+	settables []internal.Settable
+	factories []func(ctx internal.Context) internal.Future
+	batchSize int
+
+	// state
+	wg internal.WaitGroup
+}
+
+func NewBatchFuture(ctx internal.Context, batchSize int, factories []func(ctx internal.Context) internal.Future) (BatchFuture, error) {
+	var futures []internal.Future
+	var settables []internal.Settable
+	for range factories {
+		future, settable := internal.NewFuture(ctx)
+		futures = append(futures, future)
+		settables = append(settables, settable)
+	}
+
+	batchFuture := &batchFutureImpl{
+		futures:   futures,
+		settables: settables,
+		factories: factories,
+		batchSize: batchSize,
+
+		wg: internal.NewWaitGroup(ctx),
+	}
+	batchFuture.start(ctx)
+	return batchFuture, nil
+}
+
+func (b *batchFutureImpl) GetFutures() []internal.Future {
+	return b.futures
+}
+
+func (b *batchFutureImpl) start(ctx internal.Context) {
+
+	buffered := internal.NewBufferedChannel(ctx, b.batchSize) // buffered channel to limit the number of concurrent futures
+	channel := internal.NewNamedChannel(ctx, "batch-future-channel")
+	b.wg.Add(1)
+	internal.GoNamed(ctx, "batch-future-submitter", func(ctx internal.Context) {
+		defer b.wg.Done()
+
+		for i := range b.factories {
+			buffered.Send(ctx, nil)
+			channel.Send(ctx, i)
+		}
+		channel.Close()
+	})
+
+	b.wg.Add(1)
+	internal.GoNamed(ctx, "batch-future-processor", func(ctx internal.Context) {
+		defer b.wg.Done()
+
+		wgForFutures := internal.NewWaitGroup(ctx)
+
+		var idx int
+		for channel.Receive(ctx, &idx) {
+			idx := idx
+
+			wgForFutures.Add(1)
+			internal.GoNamed(ctx, "batch-future-processor-one-future", func(ctx internal.Context) {
+				defer wgForFutures.Done()
+
+				// fork a future and chain it to the processed future for user to get the result
+				f := b.factories[idx](ctx)
+				b.settables[idx].Chain(f)
+
+				// error handling is not needed here because the result is chained to the settable
+				f.Get(ctx, nil)
+				buffered.Receive(ctx, nil)
+			})
+		}
+		wgForFutures.Wait(ctx)
+	})
+}
+
+func (b *batchFutureImpl) IsReady() bool {
+	for _, future := range b.futures {
+		if !future.IsReady() {
+			return false
+		}
+	}
+	return true
+}
+
+func (b *batchFutureImpl) Get(ctx internal.Context, valuePtr interface{}) error {
+	// ensure valuePtr is a slice
+	var sliceValue reflect.Value
+	if valuePtr != nil {
+
+		switch v := reflect.ValueOf(valuePtr); v.Kind() {
+		case reflect.Ptr:
+			if v.Elem().Kind() != reflect.Slice {
+				return fmt.Errorf("valuePtr must be a pointer to a slice, got %v", v)
+			}
+			sliceValue = v.Elem()
+		case reflect.Slice:
+			sliceValue = v
+		default:
+			return fmt.Errorf("valuePtr must be a slice or a pointer to a slice, got %v", v.Kind())
+		}
+		// ensure slice size is the same as the number of futures
+		if sliceValue.Len() != len(b.futures) {
+			return fmt.Errorf("slice size must be the same as the number of futures, got %d, expected %d", sliceValue.Len(), len(b.futures))
+		}
+	}
+
+	// wait for all futures to be ready
+	b.wg.Wait(ctx)
+
+	// loop through all elements of valuePtr
+	var errs error
+	for i := range b.futures {
+		if valuePtr == nil {
+			errs = multierr.Append(errs, b.futures[i].Get(ctx, nil))
+		} else {
+			value := sliceValue.Index(i)
+			if value.Kind() != reflect.Ptr {
+				value = value.Addr()
+			}
+			// if value is nil, initialize it
+			if value.IsNil() {
+				value.Set(reflect.New(value.Type().Elem()))
+			}
+
+			e := b.futures[i].Get(ctx, value.Interface())
+			errs = multierr.Append(errs, e)
+		}
+	}
+
+	return errs
+}

internal/batch/batch_future.go

@@ -0,0 +1,218 @@
+package batch
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math/rand"
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"go.uber.org/cadence/internal"
+	"go.uber.org/cadence/testsuite"
+	"go.uber.org/multierr"
+)
+
+type batchWorkflowInput struct {
+	Concurrency int
+	TotalSize   int
+}
+
+func batchWorkflow(ctx internal.Context, input batchWorkflowInput) ([]int, error) {
+	factories := make([]func(ctx internal.Context) internal.Future, input.TotalSize)
+	for i := 0; i < input.TotalSize; i++ {
+		i := i
+		factories[i] = func(ctx internal.Context) internal.Future {
+			aCtx := internal.WithActivityOptions(ctx, internal.ActivityOptions{
+				ScheduleToStartTimeout: time.Second * 10,
+				StartToCloseTimeout:    time.Second * 10,
+			})
+			return internal.ExecuteActivity(aCtx, batchActivity, i)
+		}
+	}
+
+	batchFuture, err := NewBatchFuture(ctx, input.Concurrency, factories)
+	if err != nil {
+		return nil, err
+	}
+
+	result := make([]int, input.TotalSize)
+	err = batchFuture.Get(ctx, &result)
+	return result, err
+}
+
+func batchWorkflowUsingFutures(ctx internal.Context, input batchWorkflowInput) ([]int, error) {
+	factories := make([]func(ctx internal.Context) internal.Future, input.TotalSize)
+	for i := 0; i < input.TotalSize; i++ {
+		i := i
+		factories[i] = func(ctx internal.Context) internal.Future {
+			aCtx := internal.WithActivityOptions(ctx, internal.ActivityOptions{
+				ScheduleToStartTimeout: time.Second * 10,
+				StartToCloseTimeout:    time.Second * 10,
+			})
+			return internal.ExecuteActivity(aCtx, batchActivity, i)
+		}
+	}
+
+	batchFuture, err := NewBatchFuture(ctx, input.Concurrency, factories)
+	if err != nil {
+		return nil, err
+	}
+	result := make([]int, input.TotalSize)
+
+	for i, f := range batchFuture.GetFutures() {
+		err = f.Get(ctx, &result[i])
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return result, err
+}
+
+func batchActivity(ctx context.Context, taskID int) (int, error) {
+	select {
+	case <-ctx.Done():
+		return taskID, fmt.Errorf("batch activity %d failed: %w", taskID, ctx.Err())
+	case <-time.After(time.Duration(rand.Int63n(100))*time.Millisecond + 900*time.Millisecond):
+		return taskID, nil
+	}
+}
+
+func Test_BatchWorkflow(t *testing.T) {
+	testSuite := &testsuite.WorkflowTestSuite{}
+	env := testSuite.NewTestWorkflowEnvironment()
+
+	env.RegisterWorkflow(batchWorkflow)
+	env.RegisterActivity(batchActivity)
+
+	totalSize := 5
+	concurrency := 2
+	go func() {
+		env.ExecuteWorkflow(batchWorkflow, batchWorkflowInput{
+			Concurrency: concurrency,
+			TotalSize:   totalSize,
+		})
+	}()
+
+	// wait for maximum time it takes to complete the workflow (totalSize/concurrency) + 1 second
+	assert.Eventually(t, func() bool {
+		return env.IsWorkflowCompleted()
+	}, time.Second*time.Duration(1+float64(totalSize)/float64(concurrency)), time.Millisecond*100)
+
+	assert.Nil(t, env.GetWorkflowError())
+	var result []int
+	assert.Nil(t, env.GetWorkflowResult(&result))
+	var expected []int
+	for i := 0; i < totalSize; i++ {
+		expected = append(expected, i)
+	}
+	assert.Equal(t, expected, result)
+}
+
+func Test_BatchWorkflow_Cancel(t *testing.T) {
+	testSuite := &testsuite.WorkflowTestSuite{}
+	env := testSuite.NewTestWorkflowEnvironment()
+	env.RegisterWorkflow(batchWorkflow)
+	env.RegisterActivity(batchActivity)
+
+	totalSize := 100
+	concurrency := 10
+	go func() {
+		env.ExecuteWorkflow(batchWorkflow, batchWorkflowInput{
+			Concurrency: concurrency,
+			TotalSize:   totalSize,
+		})
+	}()
+
+	time.Sleep(time.Second*2)
+	env.CancelWorkflow()
+
+	assert.Eventually(t, func() bool {
+		return env.IsWorkflowCompleted()
+	}, time.Second*time.Duration(1+float64(totalSize)/float64(concurrency)), time.Millisecond*100)
+
+	err := env.GetWorkflowError()
+	errs := multierr.Errors(errors.Unwrap(err))
+	assert.Less(t, len(errs), totalSize, "expect at least some to succeed")
+	for _, e := range errs {
+		assert.Contains(t, e.Error(), "Canceled")
+	}
+}
+
+func Test_BatchWorkflowUsingFutures(t *testing.T) {
+	testSuite := &testsuite.WorkflowTestSuite{}
+	env := testSuite.NewTestWorkflowEnvironment()
+
+	env.RegisterWorkflow(batchWorkflowUsingFutures)
+	env.RegisterActivity(batchActivity)
+
+	totalSize := 100
+	concurrency := 20
+	go func() {
+		env.ExecuteWorkflow(batchWorkflowUsingFutures, batchWorkflowInput{
+			Concurrency: concurrency,
+			TotalSize:   totalSize,
+		})
+	}()
+
+	// wait for maximum time it takes to complete the workflow (totalSize/concurrency) + 1 second
+	assert.Eventually(t, func() bool {
+		return env.IsWorkflowCompleted()
+	}, time.Second*time.Duration(1+float64(totalSize)/float64(concurrency)), time.Millisecond*100)
+
+	assert.Nil(t, env.GetWorkflowError())
+	var result []int
+	assert.Nil(t, env.GetWorkflowResult(&result))
+	var expected []int
+	for i := 0; i < totalSize; i++ {
+		expected = append(expected, i)
+	}
+	assert.Equal(t, expected, result)
+}
+
+func futureTest(ctx internal.Context) error {
+	f, s := internal.NewFuture(ctx)
+	f2, s2 := internal.NewFuture(ctx)
+	s2.Chain(f)
+
+	wg := internal.NewWaitGroup(ctx)
+	wg.Add(1)
+	internal.GoNamed(ctx, "future-test", func(ctx internal.Context) {
+		defer wg.Done()
+		internal.Sleep(ctx, time.Second*10)
+		s.Set(1, nil)
+	})
+
+	err := f2.Get(ctx, nil)
+	if err != nil {
+		return err
+	}
+
+	err = f.Get(ctx, nil)
+	if err != nil {
+		return err
+	}
+
+	wg.Wait(ctx)
+	return err
+}
+
+func Test_Futures(t *testing.T) {
+	testSuite := &testsuite.WorkflowTestSuite{}
+	env := testSuite.NewTestWorkflowEnvironment()
+
+	env.RegisterWorkflow(futureTest)
+
+	env.ExecuteWorkflow(futureTest)
+}
+
+func Test_valuePtr(t *testing.T) {
+	slices := make([]int, 10)
+	slicePtr := &slices
+
+	fmt.Println(reflect.ValueOf(slicePtr).Elem().Len())
+}