From f277864aa5a8f59ada87ade92d71891f5a4be56f Mon Sep 17 00:00:00 2001
From: BinacsLee <bin646891055@gmail.com>
Date: Sun, 19 Sep 2021 10:09:06 +0800
Subject: [PATCH] Scheduler queue: fix calculateBackoffDuration overflow in
 extreme data cases

---
 .../internal/queue/scheduling_queue.go        |  5 ++-
 .../internal/queue/scheduling_queue_test.go   | 44 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/pkg/scheduler/internal/queue/scheduling_queue.go b/pkg/scheduler/internal/queue/scheduling_queue.go
index fc2e97b65de..8583491154a 100644
--- a/pkg/scheduler/internal/queue/scheduling_queue.go
+++ b/pkg/scheduler/internal/queue/scheduling_queue.go
@@ -742,10 +742,11 @@ func (p *PriorityQueue) getBackoffTime(podInfo *framework.QueuedPodInfo) time.Ti
 func (p *PriorityQueue) calculateBackoffDuration(podInfo *framework.QueuedPodInfo) time.Duration {
 	duration := p.podInitialBackoffDuration
 	for i := 1; i < podInfo.Attempts; i++ {
-		duration = duration * 2
-		if duration > p.podMaxBackoffDuration {
+		// Use subtraction instead of addition or multiplication to avoid overflow.
+		if duration > p.podMaxBackoffDuration-duration {
 			return p.podMaxBackoffDuration
 		}
+		duration += duration
 	}
 	return duration
 }
diff --git a/pkg/scheduler/internal/queue/scheduling_queue_test.go b/pkg/scheduler/internal/queue/scheduling_queue_test.go
index 79e1656fa0c..277a3419b82 100644
--- a/pkg/scheduler/internal/queue/scheduling_queue_test.go
+++ b/pkg/scheduler/internal/queue/scheduling_queue_test.go
@@ -19,6 +19,7 @@ package queue
 import (
 	"context"
 	"fmt"
+	"math"
 	"reflect"
 	"strings"
 	"sync"
@@ -27,10 +28,9 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
-	"k8s.io/apimachinery/pkg/runtime"
-
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/clock"
 	"k8s.io/apimachinery/pkg/util/sets"
@@ -1995,3 +1995,43 @@ func makeQueuedPodInfos(num int, timestamp time.Time) []*framework.QueuedPodInfo
 	}
 	return pInfos
 }
+
+func TestPriorityQueue_calculateBackoffDuration(t *testing.T) {
+	tests := []struct {
+		name                   string
+		initialBackoffDuration time.Duration
+		maxBackoffDuration     time.Duration
+		podInfo                *framework.QueuedPodInfo
+		want                   time.Duration
+	}{
+		{
+			name:                   "normal",
+			initialBackoffDuration: 1 * time.Nanosecond,
+			maxBackoffDuration:     32 * time.Nanosecond,
+			podInfo:                &framework.QueuedPodInfo{Attempts: 16},
+			want:                   32 * time.Nanosecond,
+		},
+		{
+			name:                   "overflow_32bit",
+			initialBackoffDuration: 1 * time.Nanosecond,
+			maxBackoffDuration:     math.MaxInt32 * time.Nanosecond,
+			podInfo:                &framework.QueuedPodInfo{Attempts: 32},
+			want:                   math.MaxInt32 * time.Nanosecond,
+		},
+		{
+			name:                   "overflow_64bit",
+			initialBackoffDuration: 1 * time.Nanosecond,
+			maxBackoffDuration:     math.MaxInt64 * time.Nanosecond,
+			podInfo:                &framework.QueuedPodInfo{Attempts: 64},
+			want:                   math.MaxInt64 * time.Nanosecond,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			q := NewTestQueue(context.Background(), newDefaultQueueSort(), WithPodInitialBackoffDuration(tt.initialBackoffDuration), WithPodMaxBackoffDuration(tt.maxBackoffDuration))
+			if got := q.calculateBackoffDuration(tt.podInfo); got != tt.want {
+				t.Errorf("PriorityQueue.calculateBackoffDuration() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}