diff --git a/internal/internal_worker.go b/internal/internal_worker.go index b743ef468..a6ad56a67 100644 --- a/internal/internal_worker.go +++ b/internal/internal_worker.go @@ -462,6 +462,33 @@ func (th *hostEnvImpl) RegisterWorkflowWithOptions( return nil } +func (th *hostEnvImpl) UnRegisterWorkflow(wf interface{}) error { + fnType := reflect.TypeOf(wf) + + var registrationName string + var functioName string + if wfName, ok := wf.(string); ok { + // we suppose the workflow name is the alias here + registrationName = wfName + } else { + if err := validateFnFormat(fnType, true); err != nil { + return err + } + functioName = getFunctionName(wf) + if alias, ok := th.getWorkflowAlias(functioName); ok { + registrationName = alias + } else { + registrationName = functioName + } + } + + th.delWorkflowFn(registrationName) + th.delWorkflowAliasByAlias(registrationName) + + return nil +} + + func (th *hostEnvImpl) RegisterActivity(af interface{}) error { return th.RegisterActivityWithOptions(af, RegisterActivityOptions{}) } @@ -498,6 +525,16 @@ func (th *hostEnvImpl) addWorkflowAlias(fnName string, alias string) { th.workflowAliasMap[fnName] = alias } +func (th *hostEnvImpl) delWorkflowAliasByAlias(alias string) { + th.Lock() + defer th.Unlock() + for k, v := range th.workflowAliasMap { + if v == alias { + delete(th.workflowAliasMap, k) + } + } +} + func (th *hostEnvImpl) getWorkflowAlias(fnName string) (string, bool) { th.Lock() defer th.Unlock() @@ -511,6 +548,13 @@ func (th *hostEnvImpl) addWorkflowFn(fnName string, wf interface{}) { th.workflowFuncMap[fnName] = wf } +func (th *hostEnvImpl) delWorkflowFn(fnName string) { + th.Lock() + defer th.Unlock() + delete(th.workflowFuncMap, fnName) +} + + func (th *hostEnvImpl) getWorkflowFn(fnName string) (interface{}, bool) { th.Lock() defer th.Unlock() diff --git a/internal/internal_worker_test.go b/internal/internal_worker_test.go index 0169f8b90..de19fbe88 100644 --- a/internal/internal_worker_test.go +++ b/internal/internal_worker_test.go @@ -521,6 +521,66 @@ func (s *internalWorkerTestSuite) TestRecordActivityHeartbeatByID() { require.NotNil(s.T(), heartbeatRequest) } +func (s *internalWorkerTestSuite) TestWorkflowUnRegistration() { + + wf1 := func(ctx Context) (string, error) { return "wf1", nil } + wf2 := func(ctx Context) (string, error) { return "wf2", nil } + + s.T().Run("Should allow re-registering an unregistered workflow", func(t *testing.T) { + RegisterWorkflow(wf1) + require.Panics(s.T(), func() { + RegisterWorkflow(wf1) + }, "should not allow to register the same workflow twice") + UnRegisterWorkflow(wf1) + require.NotPanics(s.T(), func() { + RegisterWorkflow(wf1) + }, "should allow to register the workflow once it has been unregistered") + UnRegisterWorkflow(wf1) + }) + + s.T().Run("Should not mix-up two different workflows registration/unregistration", func(t *testing.T) { + RegisterWorkflow(wf1) + RegisterWorkflow(wf2) + UnRegisterWorkflow(wf2) + require.Panics(s.T(), func() { + RegisterWorkflow(wf1) + }, "un-registration of wf2 should not affect the registration of wf1") + UnRegisterWorkflow(wf1) + }) + + s.T().Run("Should allow un-registration of workflow registered with alias", func(t *testing.T) { + RegisterWorkflowWithOptions(wf2, RegisterWorkflowOptions{Name: "alias"}) + UnRegisterWorkflow(wf2) + require.NotPanics(s.T(), func() { + RegisterWorkflowWithOptions(wf2, RegisterWorkflowOptions{Name: "alias"}) + }, "workflow should not be registered after its un-registration") + UnRegisterWorkflow(wf2) + }) + + s.T().Run("Should allow un-registration by name of workflow registered with alias", func(t *testing.T) { + RegisterWorkflowWithOptions(wf2, RegisterWorkflowOptions{Name: "alias"}) + UnRegisterWorkflow("alias") + require.NotPanics(s.T(), func() { + RegisterWorkflowWithOptions(wf2, RegisterWorkflowOptions{Name: "alias"}) + }, "workflow should not be registered after its un-registration by alias") + UnRegisterWorkflow(wf2) + }) + + s.T().Run("Should allow selective un-registration by alias name of workflow registered twice with different aliases", func(t *testing.T) { + RegisterWorkflowWithOptions(wf2, RegisterWorkflowOptions{Name: "alias1"}) + RegisterWorkflowWithOptions(wf2, RegisterWorkflowOptions{Name: "alias2"}) + UnRegisterWorkflow("alias1") + require.NotPanics(s.T(), func() { + RegisterWorkflowWithOptions(wf2, RegisterWorkflowOptions{Name: "alias1"}) + }, "workflow should not be registered after its un-registration by alias1") + require.Panics(s.T(), func() { + RegisterWorkflowWithOptions(wf2, RegisterWorkflowOptions{Name: "alias2"}) + }, "workflow alias2 still be registered after the un-registration of alias1") + UnRegisterWorkflow("alias1") + UnRegisterWorkflow("alias2") + }) +} + type activitiesCallingOptionsWorkflow struct { t *testing.T } diff --git a/internal/workflow.go b/internal/workflow.go index a264f6a52..1bcaa685a 100644 --- a/internal/workflow.go +++ b/internal/workflow.go @@ -232,6 +232,19 @@ func RegisterWorkflow(workflowFunc interface{}) { RegisterWorkflowWithOptions(workflowFunc, RegisterWorkflowOptions{}) } +// UnRegisterWorkflow - un-registers a workflow from the framework. +// The supplied workflow can be a workflow function or a string. +// In case a string is passed, that should be the alias of the workflow. +// In case a function is passed, all the aliases referring to the workflow (if any) will be unregistered as well. +// This method calls panic if workflowFunc doesn't comply with the expected format. +func UnRegisterWorkflow(wf interface{}) { + thImpl := getHostEnvironment() + err := thImpl.UnRegisterWorkflow(wf) + if err != nil { + panic(err) + } +} + // RegisterWorkflowWithOptions registers the workflow function with options // The user can use options to provide an external name for the workflow or leave it empty if no // external name is required. This can be used as