m-trieu commented on code in PR #29082:
URL: https://github.com/apache/beam/pull/29082#discussion_r1376809632


##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java:
##########
@@ -126,37 +164,40 @@ synchronized ActivateWorkResult 
activateWorkForKey(ShardedKey shardedKey, Work w
    * #activeWork}.
    */
   synchronized Optional<Work> completeWorkAndGetNextWorkForKey(
-      ShardedKey shardedKey, long workToken) {
+      ShardedKey shardedKey, WorkId workId) {
     @Nullable Queue<Work> workQueue = activeWork.get(shardedKey);
     if (workQueue == null) {
       // Work may have been completed due to clearing of stuck commits.
-      LOG.warn("Unable to complete inactive work for key {} and token {}.", 
shardedKey, workToken);
+      LOG.warn("Unable to complete inactive work for key {} and token {}.", 
shardedKey, workId);
       return Optional.empty();
     }
-    removeCompletedWorkFromQueue(workQueue, shardedKey, workToken);
+    removeCompletedWorkFromQueue(workQueue, shardedKey, workId);
     return getNextWork(workQueue, shardedKey);
   }
 
   private synchronized void removeCompletedWorkFromQueue(
-      Queue<Work> workQueue, ShardedKey shardedKey, long workToken) {
+      Queue<Work> workQueue, ShardedKey shardedKey, WorkId workId) {
     // avoid Preconditions.checkState here to prevent eagerly evaluating the
     // format string parameters for the error message.
-    Work completedWork =
+    Work workToComplete =
         Optional.ofNullable(workQueue.peek())
             .orElseThrow(
                 () ->
                     new IllegalStateException(
                         String.format(
-                            "Active key %s without work, expected token %d",
-                            shardedKey, workToken)));
+                            "Active key %s without work, expected work_token 
%d, expected cache_token %d",
+                            shardedKey, workId.workToken(), 
workId.cacheToken())));
 
-    if (completedWork.getWorkItem().getWorkToken() != workToken) {
+    if (!workToComplete.id().isForSameWork(workId) || 
workToComplete.id().isRetryOf(workId)) {

Review Comment:
   done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to