Merge locking keys into partial loading keys when using locking loadi…

…ng types (#341)
indeedeng · Sep 18, 2023 · 5410299 · 5410299
1 parent 1fe1b5f
commit 5410299
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 2 deletions.
diff --git a/integ/workflow/persistence_loading_policy/routers.go b/integ/workflow/persistence_loading_policy/routers.go
@@ -172,7 +172,7 @@ func verifyLoadedAttributes(
 	var expectedSearchAttributes []iwfidl.SearchAttribute
 	var expectedDataAttributes []iwfidl.KeyValue
 
-	if loadingType == iwfidl.ALL_WITHOUT_LOCKING {
+	if loadingType == iwfidl.ALL_WITHOUT_LOCKING || loadingType == iwfidl.PARTIAL_WITH_EXCLUSIVE_LOCK {
 		expectedSearchAttributes = []iwfidl.SearchAttribute{
 			{
 				Key:         iwfidl.PtrString(persistence.TestSearchAttributeKeywordKey),
@@ -201,7 +201,7 @@ func verifyLoadedAttributes(
 				},
 			},
 		}
-	} else if loadingType == iwfidl.PARTIAL_WITHOUT_LOCKING || loadingType == iwfidl.PARTIAL_WITH_EXCLUSIVE_LOCK {
+	} else if loadingType == iwfidl.PARTIAL_WITHOUT_LOCKING {
 		expectedSearchAttributes = []iwfidl.SearchAttribute{
 			{
 				Key:         iwfidl.PtrString(persistence.TestSearchAttributeKeywordKey),

diff --git a/service/common/utils/utils.go b/service/common/utils/utils.go
@@ -11,6 +11,24 @@ const (
 	defaultMaxApiTimeoutSeconds = 60
 )
 
+func MergeStringSlice(first, second []string) []string {
+	exists := map[string]bool{}
+	var out []string
+	for _, k := range first {
+		if !exists[k] {
+			exists[k] = true
+			out = append(out, k)
+		}
+	}
+	for _, k := range second {
+		if !exists[k] {
+			exists[k] = true
+			out = append(out, k)
+		}
+	}
+	return out
+}
+
 func MergeMap(first map[string]interface{}, second map[string]interface{}) map[string]interface{} {
 	out := make(map[string]interface{}, len(first))
 	for k, v := range first {

diff --git a/service/interpreter/persistence.go b/service/interpreter/persistence.go
@@ -4,6 +4,7 @@ import (
 	"github.com/indeedeng/iwf/gen/iwfidl"
 	"github.com/indeedeng/iwf/service"
 	"github.com/indeedeng/iwf/service/common/mapper"
+	"github.com/indeedeng/iwf/service/common/utils"
 )
 
 type PersistenceManager struct {
@@ -85,6 +86,9 @@ func (am *PersistenceManager) LoadSearchAttributes(ctx UnifiedContext, loadingPo
 	if loadingPolicy != nil {
 		loadingType = loadingPolicy.GetPersistenceLoadingType()
 		partialLoadingKeys = loadingPolicy.PartialLoadingKeys
+		if loadingType == iwfidl.PARTIAL_WITH_EXCLUSIVE_LOCK {
+			partialLoadingKeys = utils.MergeStringSlice(loadingPolicy.PartialLoadingKeys, loadingPolicy.LockingKeys)
+		}
 
 		if loadingType == iwfidl.PARTIAL_WITH_EXCLUSIVE_LOCK {
 			am.awaitAndLockForKeys(ctx, am.lockedSearchAttributeKeys, loadingPolicy.GetLockingKeys())
@@ -118,6 +122,9 @@ func (am *PersistenceManager) LoadDataObjects(ctx UnifiedContext, loadingPolicy
 	if loadingPolicy != nil {
 		loadingType = loadingPolicy.GetPersistenceLoadingType()
 		partialLoadingKeys = loadingPolicy.PartialLoadingKeys
+		if loadingType == iwfidl.PARTIAL_WITH_EXCLUSIVE_LOCK {
+			partialLoadingKeys = utils.MergeStringSlice(loadingPolicy.PartialLoadingKeys, loadingPolicy.LockingKeys)
+		}
 
 		if loadingType == iwfidl.PARTIAL_WITH_EXCLUSIVE_LOCK {
 			am.awaitAndLockForKeys(ctx, am.lockedDataObjectKeys, loadingPolicy.GetLockingKeys())