make HandleAction more robust

cadence-workflow · Jul 30, 2020 · 1ced8db · 1ced8db
1 parent 4982dae
commit 1ced8db
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 9 deletions.
diff --git a/service/history/queue/timer_queue_processor.go b/service/history/queue/timer_queue_processor.go
@@ -22,7 +22,6 @@ package queue
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"sync"
 	"sync/atomic"
@@ -309,12 +308,15 @@ func (t *timerQueueProcessor) HandleAction(clusterName string, action *Action) (
 	}
 
 	if !added {
-		return nil, errors.New("queue processor has been shutdown")
+		return nil, errProcessorShutdown
 	}
 
-	resultNotification := <-resultNotificationCh
-
-	return resultNotification.result, resultNotification.err
+	select {
+	case resultNotification := <-resultNotificationCh:
+		return resultNotification.result, resultNotification.err
+	case <-t.shutdownChan:
+		return nil, errProcessorShutdown
+	}
 }
 
 func (t *timerQueueProcessor) LockTaskProcessing() {

diff --git a/service/history/queue/transfer_queue_processor.go b/service/history/queue/transfer_queue_processor.go
@@ -56,6 +56,7 @@ const (
 
 var (
 	errUnexpectedQueueTask = errors.New("unexpected queue task")
+	errProcessorShutdown   = errors.New("queue processor has been shutdown")
 
 	maxTransferReadLevel = newTransferTaskKey(math.MaxInt64)
 )
@@ -320,12 +321,15 @@ func (t *transferQueueProcessor) HandleAction(clusterName string, action *Action
 	}
 
 	if !added {
-		return nil, errors.New("queue processor has been shutdown")
+		return nil, errProcessorShutdown
 	}
 
-	resultNotification := <-resultNotificationCh
-
-	return resultNotification.result, resultNotification.err
+	select {
+	case resultNotification := <-resultNotificationCh:
+		return resultNotification.result, resultNotification.err
+	case <-t.shutdownChan:
+		return nil, errProcessorShutdown
+	}
 }
 
 func (t *transferQueueProcessor) LockTaskProcessing() {