func param WatchOrchestrationRuntimeStatus (#47)

Update WatchOrchestrationRuntimeStatus to use a func condition param
to reduce go routine count.

Signed-off-by: joshvanl <[email protected]>

Author: JoshVanL

backend/backend.go

@@ -87,8 +87,9 @@ type Backend interface {
 
 	// WatchOrchestrationRuntimeStatus is a streaming API to watch for changes to
 	// the OrchestrtionMetadata, receiving events as and when the state changes.
+	// When the given condition is true, returns.
 	// Used over polling the metadata.
-	WatchOrchestrationRuntimeStatus(ctx context.Context, id api.InstanceID, ch chan<- *OrchestrationMetadata) error
+	WatchOrchestrationRuntimeStatus(ctx context.Context, id api.InstanceID, condition func(*OrchestrationMetadata) bool) error
 
 	// GetOrchestrationMetadata gets the metadata associated with the given orchestration instance ID.
 	//

backend/client.go

@@ -2,7 +2,6 @@ package backend
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"time"
 
@@ -15,7 +14,6 @@ import (
 	"github.com/dapr/durabletask-go/api"
 	"github.com/dapr/durabletask-go/api/helpers"
 	"github.com/dapr/durabletask-go/api/protos"
-	"github.com/dapr/kit/concurrency"
 )
 
 type TaskHubClient interface {
@@ -116,31 +114,13 @@ func (c *backendClient) WaitForOrchestrationCompletion(ctx context.Context, id a
 }
 
 func (c *backendClient) waitForOrchestrationCondition(ctx context.Context, id api.InstanceID, condition func(metadata *OrchestrationMetadata) bool) (*OrchestrationMetadata, error) {
-	ch := make(chan *protos.OrchestrationMetadata)
 	var metadata *protos.OrchestrationMetadata
-	err := concurrency.NewRunnerManager(
-		func(ctx context.Context) error {
-			return c.be.WatchOrchestrationRuntimeStatus(ctx, id, ch)
-		},
-		func(ctx context.Context) error {
-			for {
-				select {
-				case <-ctx.Done():
-					return ctx.Err()
-				case metadata = <-ch:
-					if condition(metadata) {
-						return nil
-					}
-				}
-			}
-		},
-	).Run(ctx)
-
-	if err != nil || ctx.Err() != nil {
-		return nil, errors.Join(err, ctx.Err())
-	}
+	err := c.be.WatchOrchestrationRuntimeStatus(ctx, id, func(m *OrchestrationMetadata) bool {
+		metadata = m
+		return condition(m)
+	})
 
-	return metadata, nil
+	return metadata, err
 }
 
 // TerminateOrchestration enqueues a message to terminate a running orchestration, causing it to stop receiving new events and

backend/executor.go

@@ -21,7 +21,6 @@ import (
 	"github.com/dapr/durabletask-go/api"
 	"github.com/dapr/durabletask-go/api/helpers"
 	"github.com/dapr/durabletask-go/api/protos"
-	"github.com/dapr/kit/concurrency"
 )
 
 var emptyCompleteTaskResponse = &protos.CompleteTaskResponse{}
@@ -622,28 +621,13 @@ func (g *grpcExecutor) WaitForInstanceStart(ctx context.Context, req *protos.Get
 func (g *grpcExecutor) waitForInstance(ctx context.Context, req *protos.GetInstanceRequest, condition func(*OrchestrationMetadata) bool) (*protos.GetInstanceResponse, error) {
 	iid := api.InstanceID(req.InstanceId)
 
-	ch := make(chan *protos.OrchestrationMetadata)
 	var metadata *protos.OrchestrationMetadata
-	err := concurrency.NewRunnerManager(
-		func(ctx context.Context) error {
-			return g.backend.WatchOrchestrationRuntimeStatus(ctx, iid, ch)
-		},
-		func(ctx context.Context) error {
-			for {
-				select {
-				case <-ctx.Done():
-					return ctx.Err()
-				case metadata = <-ch:
-					if condition(metadata) {
-						return nil
-					}
-				}
-			}
-		},
-	).Run(ctx)
-
-	if err != nil || ctx.Err() != nil {
-		return nil, errors.Join(err, ctx.Err())
+	err := g.backend.WatchOrchestrationRuntimeStatus(ctx, iid, func(m *OrchestrationMetadata) bool {
+		metadata = m
+		return condition(m)
+	})
+	if err != nil {
+		return nil, err
 	}
 
 	if metadata == nil {

backend/postgres/postgres.go

@@ -155,7 +155,7 @@ func (be *postgresBackend) NextOrchestrationWorkItem(ctx context.Context) (*back
 	}
 }
 
-func (be *postgresBackend) WatchOrchestrationRuntimeStatus(ctx context.Context, id api.InstanceID, ch chan<- *backend.OrchestrationMetadata) error {
+func (be *postgresBackend) WatchOrchestrationRuntimeStatus(ctx context.Context, id api.InstanceID, fn func(*backend.OrchestrationMetadata) bool) error {
 	b := backoff.ExponentialBackOff{
 		InitialInterval:     100 * time.Millisecond,
 		MaxInterval:         10 * time.Second,
@@ -181,10 +181,8 @@ func (be *postgresBackend) WatchOrchestrationRuntimeStatus(ctx context.Context,
 				return err
 			}
 
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			case ch <- meta:
+			if fn(meta) {
+				return nil
 			}
 		}
 	}

backend/sqlite/sqlite.go

@@ -644,7 +644,7 @@ func (be *sqliteBackend) AddNewOrchestrationEvent(ctx context.Context, iid api.I
 	return nil
 }
 
-func (be *sqliteBackend) WatchOrchestrationRuntimeStatus(ctx context.Context, id api.InstanceID, ch chan<- *backend.OrchestrationMetadata) error {
+func (be *sqliteBackend) WatchOrchestrationRuntimeStatus(ctx context.Context, id api.InstanceID, fn func(*backend.OrchestrationMetadata) bool) error {
 	b := backoff.ExponentialBackOff{
 		InitialInterval:     100 * time.Millisecond,
 		MaxInterval:         10 * time.Second,
@@ -670,10 +670,8 @@ func (be *sqliteBackend) WatchOrchestrationRuntimeStatus(ctx context.Context, id
 				return err
 			}
 
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			case ch <- meta:
+			if fn(meta) {
+				return nil
 			}
 		}
 	}