fix selector (#416)

* fix selector

* add unit test for selector

* gofmt

Author: yiminc

internal/internal_coroutines_test.go

@@ -353,6 +353,51 @@ func TestBlockingSelectAsyncSend(t *testing.T) {
 	require.EqualValues(t, expected, history)
 }
 
+func TestBlockingSelectAsyncSend2(t *testing.T) {
+	var history []string
+	d := newDispatcher(background, func(ctx Context) {
+		c1 := NewBufferedChannel(ctx, 100)
+		c2 := NewBufferedChannel(ctx, 100)
+		s := NewSelector(ctx)
+		s.
+			AddReceive(c1, func(c Channel, more bool) {
+				assert.True(t, more)
+				var v string
+				c.Receive(ctx, &v)
+				history = append(history, fmt.Sprintf("c1-%v", v))
+			}).
+			AddReceive(c2, func(c Channel, more bool) {
+				assert.True(t, more)
+				var v string
+				c.Receive(ctx, &v)
+				history = append(history, fmt.Sprintf("c2-%v", v))
+			})
+
+		history = append(history, "send-s2")
+		c2.SendAsync("s2")
+		history = append(history, "select-0")
+		s.Select(ctx)
+		history = append(history, "send-s1")
+		c1.SendAsync("s1")
+		history = append(history, "select-1")
+		s.Select(ctx)
+		history = append(history, "done")
+	})
+	d.ExecuteUntilAllBlocked()
+	require.True(t, d.IsDone(), strings.Join(history, "\n"))
+
+	expected := []string{
+		"send-s2",
+		"select-0",
+		"c2-s2",
+		"send-s1",
+		"select-1",
+		"c1-s1",
+		"done",
+	}
+	require.EqualValues(t, expected, history)
+}
+
 func TestSendSelect(t *testing.T) {
 	var history []string
 	d := newDispatcher(background, func(ctx Context) {

internal/internal_workflow.go

@@ -880,6 +880,11 @@ func (s *selectorImpl) Select(ctx Context) {
 			}
 			v, ok, more := pair.channel.receiveAsyncImpl(callback)
 			if ok || !more {
+				// Select() returns in this case/branch. The callback won't be called for this case. However, callback
+				// will be called for previous cases/branches. We should set readyBranch so that when other case/branch
+				// become ready they won't consume the value for this Select() call.
+				readyBranch = func() {
+				}
 				c.recValue = &v
 				f(c, more)
 				return
@@ -900,6 +905,11 @@ func (s *selectorImpl) Select(ctx Context) {
 			}
 			ok := pair.channel.sendAsyncImpl(*pair.sendValue, p)
 			if ok {
+				// Select() returns in this case/branch. The callback won't be called for this case. However, callback
+				// will be called for previous cases/branches. We should set readyBranch so that when other case/branch
+				// become ready they won't consume the value for this Select() call.
+				readyBranch = func() {
+				}
 				f()
 				return
 			}
@@ -918,6 +928,11 @@ func (s *selectorImpl) Select(ctx Context) {
 			}
 			_, ok, _ := p.future.GetAsync(callback)
 			if ok {
+				// Select() returns in this case/branch. The callback won't be called for this case. However, callback
+				// will be called for previous cases/branches. We should set readyBranch so that when other case/branch
+				// become ready they won't consume the value for this Select() call.
+				readyBranch = func() {
+				}
 				p.futureFunc = nil
 				f(p.future)
 				return

internal/internal_workflow_testsuite_test.go

@@ -1536,3 +1536,69 @@ func (s *WorkflowTestSuiteUnitTest) Test_WorkflowIDReusePolicy() {
 	s.NoError(env.GetWorkflowResult(&actualResult))
 	s.Equal("hello_world", actualResult)
 }
+
+func (s *WorkflowTestSuiteUnitTest) Test_Channel() {
+	workflowFn := func(ctx Context) error {
+
+		signalCh := GetSignalChannel(ctx, "test-signal")
+		doneCh := NewBufferedChannel(ctx, 100)
+		selector := NewSelector(ctx)
+
+		selector.AddReceive(signalCh, func(c Channel, more bool) {
+		}).AddReceive(doneCh, func(c Channel, more bool) {
+			var doneSignal string
+			c.Receive(ctx, &doneSignal)
+		})
+
+		fanoutChs := []Channel{NewBufferedChannel(ctx, 100), NewBufferedChannel(ctx, 100)}
+
+		processedCount := 0
+		runningCount := 0
+
+	mainLoop:
+		for {
+			selector.Select(ctx)
+			var signal string
+			if !signalCh.ReceiveAsync(&signal) {
+				if runningCount > 0 {
+					continue mainLoop
+				}
+
+				if processedCount < 4 {
+					continue mainLoop
+				}
+
+				// continue as new
+				return NewContinueAsNewError(ctx, "this-workflow-fn")
+			}
+
+			for i := range fanoutChs {
+				ch := fanoutChs[i]
+				ch.SendAsync(signal)
+				processedCount++
+				runningCount++
+				Go(ctx, func(ctx Context) {
+					doneCh.SendAsync("done")
+					runningCount--
+				})
+			}
+		}
+
+		return nil
+	}
+
+	RegisterWorkflow(workflowFn)
+	env := s.NewTestWorkflowEnvironment()
+
+	env.RegisterDelayedCallback(func() {
+		env.SignalWorkflow("test-signal", "s1")
+		env.SignalWorkflow("test-signal", "s2")
+	}, time.Minute)
+
+	env.ExecuteWorkflow(workflowFn)
+
+	s.True(env.IsWorkflowCompleted())
+	s.Error(env.GetWorkflowError())
+	_, ok := env.GetWorkflowError().(*ContinueAsNewError)
+	s.True(ok)
+}