Fix overflow

bitfinity-network · Jun 10, 2024 · ea93f0f · ea93f0f
1 parent 567b544
commit ea93f0f
Showing 1 changed file with 25 additions and 1 deletion.
diff --git a/ic-task-scheduler/src/retry.rs b/ic-task-scheduler/src/retry.rs
@@ -49,7 +49,7 @@ impl RetryPolicy {
             match self {
                 RetryPolicy::None => false,
                 RetryPolicy::Infinite => true,
-                RetryPolicy::MaxRetries { retries: attempts } => *attempts + 1 > failed_attempts,
+                RetryPolicy::MaxRetries { retries: attempts } => *attempts >= failed_attempts,
             }
         }
     }
@@ -329,4 +329,28 @@ pub mod test {
         assert_eq!((true, 34), retry_strategy.should_retry(1));
         assert_eq!((false, 34), retry_strategy.should_retry(2));
     }
+
+    #[test]
+    fn retry_policy_hedge_cases() {
+        assert!(RetryPolicy::None.should_retry(0));
+        assert!(!RetryPolicy::None.should_retry(1));
+        assert!(!RetryPolicy::None.should_retry(u32::MAX));
+        assert!(!RetryPolicy::None.should_retry(u32::MAX-1));
+
+        assert!(RetryPolicy::MaxRetries { retries: 0 }.should_retry(0));
+        assert!(!RetryPolicy::MaxRetries { retries: 0 }.should_retry(1));
+        assert!(!RetryPolicy::MaxRetries { retries: 0 }.should_retry(u32::MAX-1));
+        assert!(!RetryPolicy::MaxRetries { retries: 0 }.should_retry(u32::MAX));
+
+        assert!(RetryPolicy::MaxRetries { retries: u32::MAX }.should_retry(0));
+        assert!(RetryPolicy::MaxRetries { retries: u32::MAX }.should_retry(1));
+        assert!(RetryPolicy::MaxRetries { retries: u32::MAX }.should_retry(u32::MAX-1));
+        assert!(RetryPolicy::MaxRetries { retries: u32::MAX }.should_retry(u32::MAX));
+
+        assert!(RetryPolicy::Infinite.should_retry(0));
+        assert!(RetryPolicy::Infinite.should_retry(1));
+        assert!(RetryPolicy::Infinite.should_retry(u32::MAX));
+        assert!(RetryPolicy::Infinite.should_retry(u32::MAX-1));
+
+    }
 }