Use timer instead of getting deadline from context

neil-xie · Jan 13, 2025 · 46fd4c3 · 46fd4c3
1 parent ac181bc
commit 46fd4c3
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 13 deletions.
diff --git a/service/matching/handler/engine.go b/service/matching/handler/engine.go
@@ -539,10 +539,7 @@ pollLoop:
 		if err != nil {
 			return nil, fmt.Errorf("couldn't load tasklist namanger: %w", err)
 		}
-		pollerWaitTimeMs := int64(0)
-		if deadline, ok := pollerCtx.Deadline(); ok {
-			pollerWaitTimeMs = int64(time.Until(deadline).Milliseconds())
-		}
+		startT := time.Now() // Record the start time
 		task, err := tlMgr.GetTask(pollerCtx, nil)
 		if err != nil {
 			// TODO: Is empty poll the best reply for errPumpClosed?
@@ -570,7 +567,7 @@ pollLoop:
 					LoadBalancerHints: tlMgr.LoadBalancerHints(),
 					AutoConfigHint: &types.AutoConfigHint{
 						EnableAutoConfig:   tlMgr.EnableClientAutoConfig(),
-						PollerWaitTimeInMs: pollerWaitTimeMs,
+						PollerWaitTimeInMs: time.Since(startT).Milliseconds(),
 					},
 				}, nil
 			}
@@ -730,10 +727,7 @@ pollLoop:
 		if err != nil {
 			return nil, fmt.Errorf("couldn't load tasklist namanger: %w", err)
 		}
-		pollerWaitTimeMs := int64(0)
-		if deadline, ok := pollerCtx.Deadline(); ok {
-			pollerWaitTimeMs = int64(time.Until(deadline).Milliseconds())
-		}
+		startT := time.Now() // Record the start time
 		task, err := tlMgr.GetTask(pollerCtx, maxDispatch)
 		if err != nil {
 			// TODO: Is empty poll the best reply for errPumpClosed?
@@ -743,7 +737,7 @@ pollLoop:
 					LoadBalancerHints: tlMgr.LoadBalancerHints(),
 					AutoConfigHint: &types.AutoConfigHint{
 						EnableAutoConfig:   tlMgr.EnableClientAutoConfig(),
-						PollerWaitTimeInMs: pollerWaitTimeMs,
+						PollerWaitTimeInMs: time.Since(startT).Milliseconds(),
 					},
 				}, nil
 			}

diff --git a/service/matching/handler/engine_integration_test.go b/service/matching/handler/engine_integration_test.go
@@ -311,7 +311,12 @@ func (s *matchingEngineSuite) PollForDecisionTasksResultTest() {
 	}
 	resp, err := pollTask(s.matchingEngine, s.handlerContext, pollReq)
 	s.NoError(err)
-	s.Equal(&pollTaskResponse{}, resp)
+	s.Equal(&pollTaskResponse{
+		AutoConfigHint: &types.AutoConfigHint{
+			EnableAutoConfig:   false,
+			PollerWaitTimeInMs: 0,
+		},
+	}, resp)
 	// add task to sticky tasklist again, this time it should pass
 	_, err = addTask(s.matchingEngine, s.handlerContext, addRequest)
 	s.NoError(err)
@@ -364,7 +369,12 @@ func (s *matchingEngineSuite) PollForTasksEmptyResultTest(callContext context.Co
 		}
 		pollResp, err := pollTask(s.matchingEngine, s.handlerContext, pollReq)
 		s.NoError(err)
-		s.Equal(&pollTaskResponse{}, pollResp)
+		s.Equal(&pollTaskResponse{
+			AutoConfigHint: &types.AutoConfigHint{
+				EnableAutoConfig:   false,
+				PollerWaitTimeInMs: 0,
+			},
+		}, pollResp)
 
 		if taskType == persistence.TaskListTypeActivity {
 			taskListType = types.TaskListTypeActivity
@@ -952,7 +962,12 @@ func (s *matchingEngineSuite) PollWithExpiredContext(taskType int) {
 	s.handlerContext.Context = ctx
 	resp, err := pollTask(s.matchingEngine, s.handlerContext, pollReq)
 	s.Nil(err)
-	s.Equal(&pollTaskResponse{}, resp)
+	s.Equal(&pollTaskResponse{
+		AutoConfigHint: &types.AutoConfigHint{
+			EnableAutoConfig:   false,
+			PollerWaitTimeInMs: 0,
+		},
+	}, resp)
 }
 
 func (s *matchingEngineSuite) TestMultipleEnginesActivitiesRangeStealing() {