temporal/go/common/mutex.go - buildbot - Git at Google

 package common

 import (
 	"context"
 	"fmt"
 	"time"

 	"go.temporal.io/sdk/activity"
 	"go.temporal.io/sdk/client"
 	"go.temporal.io/sdk/temporal"
 	"go.temporal.io/sdk/worker"
 	"go.temporal.io/sdk/workflow"
 )

 const (
 	// AcquireLockSignalName signal channel name for lock acquisition
 	AcquireLockSignalName = "acquire-lock-event"
 	// RequestLockSignalName channel name for request lock
 	RequestLockSignalName = "request-lock-event"
 	// MutexLockWorkflowName is the name of the temporal workflow mutexWorkflow
 	mutexLockWorkflowName = "perf.mutex-lock"
 )

 type contextKey struct{}

 var clientContextKey = &contextKey{}

 type UnlockFunc func() error

 // Mutex is a struct used by temporal to lock resources to allow workflows
 // to utilize a resource sequentially. Mutex workflows hold onto swarming bots
 // to ensure that they execute pairwise tasks in immediate sequential order
 // without being interrupted by swarming tasks from other jobs.
 // Mutex should only be used by run_benchmark or workflows that trigger
 // run_benchmark
 type Mutex struct {
 	// currentWorkflowID is the temporal workflow that needs the resource
 	currentWorkflowID string
 	// lockNamespace is a namespace for the mutex workflow
 	lockNamespace string
 }

 // NewMutex initializes mutex.
 // currentWorkflowID is the temporal workflow that needs the resource.
 // lockNamespace is a namespace for the mutex workflow.
 func NewMutex(currentWorkflowID string, lockNamespace string) *Mutex {
 	return &Mutex{
 		currentWorkflowID: currentWorkflowID,
 		lockNamespace:     lockNamespace,
 	}
 }

 // RegisterTemporal registers the temporal workflows and activities
 // that are needed to run Lock()
 func RegisterTemporal(w worker.Worker) *worker.Worker {
 	w.RegisterActivity(signalWithStartMutexWorkflowActivity)
 	w.RegisterWorkflowWithOptions(mutexWorkflow, workflow.RegisterOptions{Name: mutexLockWorkflowName})
 	return &w
 }

 // Lock applies the lock on resourceID by triggering a temporal workflow
 // that creates a workflow with the resourceID as part of the workflowID.
 // All external workflows outside of the mutex namespace that require
 // resourceID will have to wait for the resource until the lock is released
 // via UnlockFunc.
 // The resourceID can be any arbitrary string.
 // Example usage:
 // func SampleWorkflowWithMutex(resourceID):
 //
 //	m := NewMutex("workflow", "namespace")
 //	unlockFunc, err := m.Lock(ctx, resourceID, 10*time.Minute)
 //	< perform workflow actions >
 //	_ = unlockFunc() // this releases the lock on the resource
 //	< remaining workflow actions >
 //
 // run a bunch of workflows in parallel, but because each workflow requires
 // the same resource, they will not start until the lock is released
 // func WorkflowThatNeedsSpecificResource:
 //
 //	rc := workflow.NewBufferedChannel(ctx, numChildren)
 //	ec := workflow.NewBufferedChannel(ctx, numChildren)
 //	wg := workflow.NewWaitGroup(ctx)
 //	wg.Add(numChildren)
 //	for numChildren {
 //		workflow.Go(ctx, func(gCtx workflow.Context) {
 //			defer wg.Done
 //			if err := workflow.ExecuteChildWorkflow(gCtx, SampleWorkflowWithMutex, "bot-123").Get(); err != nil {
 //				ec.Send(gCtx, err)
 //				return
 //			}
 //			rc.Send(gCtx, status)
 //		})
 //
 // Also, if another workflow starts and calls SampleWorkflowWithMutex("bot-123")
 // while `WorkflowThatNeedsSpecificResource` is still running, it, will also wait until
 // the lock on "bot-123" is released via unlockFunc().
 func (s *Mutex) Lock(ctx workflow.Context, resourceID string, unlockTimeout time.Duration) (UnlockFunc, error) {

 	activityCtx := workflow.WithLocalActivityOptions(ctx, workflow.LocalActivityOptions{
 		ScheduleToCloseTimeout: time.Minute,
 		RetryPolicy: &temporal.RetryPolicy{
 			InitialInterval:    time.Second,
 			BackoffCoefficient: 2.0,
 			MaximumInterval:    time.Minute,
 			MaximumAttempts:    5,
 		},
 	})

 	var execution workflow.Execution
 	err := workflow.ExecuteLocalActivity(activityCtx, signalWithStartMutexWorkflowActivity, s.lockNamespace, resourceID, s.currentWorkflowID, unlockTimeout).Get(ctx, &execution)
 	if err != nil {
 		return nil, err
 	}

 	var releaseLockChannelName string
 	workflow.GetSignalChannel(ctx, AcquireLockSignalName).
 		Receive(ctx, &releaseLockChannelName)

 	unlockFunc := func() error {
 		return workflow.SignalExternalWorkflow(ctx, execution.ID, execution.RunID,
 			releaseLockChannelName, "releaseLock").Get(ctx, nil)
 	}
 	return unlockFunc, nil
 }

 // mutexWorkflow is a temporal workflow that verifies if the lock if available and
 // will poll the lock's availability until it is released or times out (unlockTimeout)
 // It is meant to only be triggered by signalWithStartMutexWorkflowActivity.
 // Note the resourceID is appended to the workflowID when it is triggered by
 // signalWithStartMutexWorkflowActivity, but is not used in this workflow.
 // Temporal will still track the workflow inputs and outputs, so including the resourceID
 // is for debugging convenience on the temporal UI.
 func mutexWorkflow(ctx workflow.Context, namespace string, resourceID string, unlockTimeout time.Duration) error {
 	logger := workflow.GetLogger(ctx)
 	logger.Info("started", "currentWorkflowID", workflow.GetInfo(ctx).WorkflowExecution.ID)
 	var ack string
 	requestLockCh := workflow.GetSignalChannel(ctx, RequestLockSignalName)
 	for {
 		var senderWorkflowID string
 		// Check if the lock is used by another workflow
 		// If not, break out of the loop
 		if !requestLockCh.ReceiveAsync(&senderWorkflowID) {
 			logger.Info("no more signals")
 			break
 		}
 		var releaseLockChannelName string
 		_ = workflow.SideEffect(ctx, func(ctx workflow.Context) interface{} {
 			return generateUnlockChannelName(senderWorkflowID)
 		}).Get(&releaseLockChannelName)
 		logger.Info("generated release lock channel name", "releaseLockChannelName", releaseLockChannelName)
 		// Send release lock channel name back to another workflow holding the lock, via senderWorkflowID,
 		// so that it can release the lock using release lock channel name
 		// .Get(ctx, nil) blocks until the signal is sent.
 		if err := workflow.SignalExternalWorkflow(ctx, senderWorkflowID, "", AcquireLockSignalName, releaseLockChannelName).Get(ctx, nil); err != nil {
 			// If the external workflow at senderWorkflowID is closed (terminated/canceled/timeouted/completed/etc)
 			// this would return error. Instead of failing the mutex workflow, release the lock immediately.
 			// Mutex workflow failing would lead to all workflows that have sent requestLock will be waiting.
 			logger.Info("SignalExternalWorkflow error", "Error", err)
 			continue
 		}
 		logger.Info("signaled external workflow")
 		selector := workflow.NewSelector(ctx)
 		selector.AddFuture(workflow.NewTimer(ctx, unlockTimeout), func(f workflow.Future) {
 			logger.Info("unlockTimeout exceeded")
 		})
 		selector.AddReceive(workflow.GetSignalChannel(ctx, releaseLockChannelName), func(c workflow.ReceiveChannel, more bool) {
 			c.Receive(ctx, &ack)
 			logger.Info("release signal received")
 		})
 		selector.Select(ctx)
 	}
 	return nil
 }

 // signalWithStartMutexWorkflowActivity triggers the MutexWorkflow using
 // the namespace and resourceID as a unique workflow to prevent other
 // workflows from accessing the same resourceID.
 func signalWithStartMutexWorkflowActivity(ctx context.Context, namespace string, resourceID string, senderWorkflowID string, unlockTimeout time.Duration) (*workflow.Execution, error) {

 	c := ctx.Value(clientContextKey).(client.Client)
 	workflowID := fmt.Sprintf(
 		"%s:%s:%s",
 		"mutex",
 		namespace,
 		resourceID,
 	)
 	workflowOptions := client.StartWorkflowOptions{
 		ID:        workflowID,
 		TaskQueue: "mutex",
 		RetryPolicy: &temporal.RetryPolicy{
 			InitialInterval:    time.Second,
 			BackoffCoefficient: 2.0,
 			MaximumInterval:    time.Minute,
 			MaximumAttempts:    5,
 		},
 	}
 	wr, err := c.SignalWithStartWorkflow(ctx, workflowID, RequestLockSignalName, senderWorkflowID, workflowOptions, mutexWorkflow, namespace, resourceID, unlockTimeout)

 	if err != nil {
 		activity.GetLogger(ctx).Error("Unable to signal with start workflow", "Error", err)
 	} else {
 		activity.GetLogger(ctx).Info("Signaled and started Workflow", "WorkflowID", wr.GetID(), "RunID", wr.GetRunID())
 	}

 	return &workflow.Execution{
 		ID:    wr.GetID(),
 		RunID: wr.GetRunID(),
 	}, nil
 }

 // generateUnlockChannelName generates release lock channel name
 func generateUnlockChannelName(senderWorkflowID string) string {
 	return fmt.Sprintf("unlock-event-%s", senderWorkflowID)
 }
	package common

	import (
	"context"
	"fmt"
	"time"

	"go.temporal.io/sdk/activity"
	"go.temporal.io/sdk/client"
	"go.temporal.io/sdk/temporal"
	"go.temporal.io/sdk/worker"
	"go.temporal.io/sdk/workflow"
	)

	const (
	// AcquireLockSignalName signal channel name for lock acquisition
	AcquireLockSignalName = "acquire-lock-event"
	// RequestLockSignalName channel name for request lock
	RequestLockSignalName = "request-lock-event"
	// MutexLockWorkflowName is the name of the temporal workflow mutexWorkflow
	mutexLockWorkflowName = "perf.mutex-lock"
	)

	type contextKey struct{}

	var clientContextKey = &contextKey{}

	type UnlockFunc func() error

	// Mutex is a struct used by temporal to lock resources to allow workflows
	// to utilize a resource sequentially. Mutex workflows hold onto swarming bots
	// to ensure that they execute pairwise tasks in immediate sequential order
	// without being interrupted by swarming tasks from other jobs.
	// Mutex should only be used by run_benchmark or workflows that trigger
	// run_benchmark
	type Mutex struct {
	// currentWorkflowID is the temporal workflow that needs the resource
	currentWorkflowID string
	// lockNamespace is a namespace for the mutex workflow
	lockNamespace string
	}

	// NewMutex initializes mutex.
	// currentWorkflowID is the temporal workflow that needs the resource.
	// lockNamespace is a namespace for the mutex workflow.
	func NewMutex(currentWorkflowID string, lockNamespace string) *Mutex {
	return &Mutex{
	currentWorkflowID: currentWorkflowID,
	lockNamespace: lockNamespace,
	}
	}

	// RegisterTemporal registers the temporal workflows and activities
	// that are needed to run Lock()
	func RegisterTemporal(w worker.Worker) *worker.Worker {
	w.RegisterActivity(signalWithStartMutexWorkflowActivity)
	w.RegisterWorkflowWithOptions(mutexWorkflow, workflow.RegisterOptions{Name: mutexLockWorkflowName})
	return &w
	}

	// Lock applies the lock on resourceID by triggering a temporal workflow
	// that creates a workflow with the resourceID as part of the workflowID.
	// All external workflows outside of the mutex namespace that require
	// resourceID will have to wait for the resource until the lock is released
	// via UnlockFunc.
	// The resourceID can be any arbitrary string.
	// Example usage:
	// func SampleWorkflowWithMutex(resourceID):
	//
	// m := NewMutex("workflow", "namespace")
	// unlockFunc, err := m.Lock(ctx, resourceID, 10*time.Minute)
	// < perform workflow actions >
	// _ = unlockFunc() // this releases the lock on the resource
	// < remaining workflow actions >
	//
	// run a bunch of workflows in parallel, but because each workflow requires
	// the same resource, they will not start until the lock is released
	// func WorkflowThatNeedsSpecificResource:
	//
	// rc := workflow.NewBufferedChannel(ctx, numChildren)
	// ec := workflow.NewBufferedChannel(ctx, numChildren)
	// wg := workflow.NewWaitGroup(ctx)
	// wg.Add(numChildren)
	// for numChildren {
	// workflow.Go(ctx, func(gCtx workflow.Context) {
	// defer wg.Done
	// if err := workflow.ExecuteChildWorkflow(gCtx, SampleWorkflowWithMutex, "bot-123").Get(); err != nil {
	// ec.Send(gCtx, err)
	// return
	// }
	// rc.Send(gCtx, status)
	// })
	//
	// Also, if another workflow starts and calls SampleWorkflowWithMutex("bot-123")
	// while `WorkflowThatNeedsSpecificResource` is still running, it, will also wait until
	// the lock on "bot-123" is released via unlockFunc().
	func (s *Mutex) Lock(ctx workflow.Context, resourceID string, unlockTimeout time.Duration) (UnlockFunc, error) {

	activityCtx := workflow.WithLocalActivityOptions(ctx, workflow.LocalActivityOptions{
	ScheduleToCloseTimeout: time.Minute,
	RetryPolicy: &temporal.RetryPolicy{
	InitialInterval: time.Second,
	BackoffCoefficient: 2.0,
	MaximumInterval: time.Minute,
	MaximumAttempts: 5,
	},
	})

	var execution workflow.Execution
	err := workflow.ExecuteLocalActivity(activityCtx, signalWithStartMutexWorkflowActivity, s.lockNamespace, resourceID, s.currentWorkflowID, unlockTimeout).Get(ctx, &execution)
	if err != nil {
	return nil, err
	}

	var releaseLockChannelName string
	workflow.GetSignalChannel(ctx, AcquireLockSignalName).
	Receive(ctx, &releaseLockChannelName)

	unlockFunc := func() error {
	return workflow.SignalExternalWorkflow(ctx, execution.ID, execution.RunID,
	releaseLockChannelName, "releaseLock").Get(ctx, nil)
	}
	return unlockFunc, nil
	}

	// mutexWorkflow is a temporal workflow that verifies if the lock if available and
	// will poll the lock's availability until it is released or times out (unlockTimeout)
	// It is meant to only be triggered by signalWithStartMutexWorkflowActivity.
	// Note the resourceID is appended to the workflowID when it is triggered by
	// signalWithStartMutexWorkflowActivity, but is not used in this workflow.
	// Temporal will still track the workflow inputs and outputs, so including the resourceID
	// is for debugging convenience on the temporal UI.
	func mutexWorkflow(ctx workflow.Context, namespace string, resourceID string, unlockTimeout time.Duration) error {
	logger := workflow.GetLogger(ctx)
	logger.Info("started", "currentWorkflowID", workflow.GetInfo(ctx).WorkflowExecution.ID)
	var ack string
	requestLockCh := workflow.GetSignalChannel(ctx, RequestLockSignalName)
	for {
	var senderWorkflowID string
	// Check if the lock is used by another workflow
	// If not, break out of the loop
	if !requestLockCh.ReceiveAsync(&senderWorkflowID) {
	logger.Info("no more signals")
	break
	}
	var releaseLockChannelName string
	_ = workflow.SideEffect(ctx, func(ctx workflow.Context) interface{} {
	return generateUnlockChannelName(senderWorkflowID)
	}).Get(&releaseLockChannelName)
	logger.Info("generated release lock channel name", "releaseLockChannelName", releaseLockChannelName)
	// Send release lock channel name back to another workflow holding the lock, via senderWorkflowID,
	// so that it can release the lock using release lock channel name
	// .Get(ctx, nil) blocks until the signal is sent.
	if err := workflow.SignalExternalWorkflow(ctx, senderWorkflowID, "", AcquireLockSignalName, releaseLockChannelName).Get(ctx, nil); err != nil {
	// If the external workflow at senderWorkflowID is closed (terminated/canceled/timeouted/completed/etc)
	// this would return error. Instead of failing the mutex workflow, release the lock immediately.
	// Mutex workflow failing would lead to all workflows that have sent requestLock will be waiting.
	logger.Info("SignalExternalWorkflow error", "Error", err)
	continue
	}
	logger.Info("signaled external workflow")
	selector := workflow.NewSelector(ctx)
	selector.AddFuture(workflow.NewTimer(ctx, unlockTimeout), func(f workflow.Future) {
	logger.Info("unlockTimeout exceeded")
	})
	selector.AddReceive(workflow.GetSignalChannel(ctx, releaseLockChannelName), func(c workflow.ReceiveChannel, more bool) {
	c.Receive(ctx, &ack)
	logger.Info("release signal received")
	})
	selector.Select(ctx)
	}
	return nil
	}

	// signalWithStartMutexWorkflowActivity triggers the MutexWorkflow using
	// the namespace and resourceID as a unique workflow to prevent other
	// workflows from accessing the same resourceID.
	func signalWithStartMutexWorkflowActivity(ctx context.Context, namespace string, resourceID string, senderWorkflowID string, unlockTimeout time.Duration) (*workflow.Execution, error) {

	c := ctx.Value(clientContextKey).(client.Client)
	workflowID := fmt.Sprintf(
	"%s:%s:%s",
	"mutex",
	namespace,
	resourceID,
	)
	workflowOptions := client.StartWorkflowOptions{
	ID: workflowID,
	TaskQueue: "mutex",
	RetryPolicy: &temporal.RetryPolicy{
	InitialInterval: time.Second,
	BackoffCoefficient: 2.0,
	MaximumInterval: time.Minute,
	MaximumAttempts: 5,
	},
	}
	wr, err := c.SignalWithStartWorkflow(ctx, workflowID, RequestLockSignalName, senderWorkflowID, workflowOptions, mutexWorkflow, namespace, resourceID, unlockTimeout)

	if err != nil {
	activity.GetLogger(ctx).Error("Unable to signal with start workflow", "Error", err)
	} else {
	activity.GetLogger(ctx).Info("Signaled and started Workflow", "WorkflowID", wr.GetID(), "RunID", wr.GetRunID())
	}

	return &workflow.Execution{
	ID: wr.GetID(),
	RunID: wr.GetRunID(),
	}, nil
	}

	// generateUnlockChannelName generates release lock channel name
	func generateUnlockChannelName(senderWorkflowID string) string {
	return fmt.Sprintf("unlock-event-%s", senderWorkflowID)
	}