From 46253aa78ec71cda303444b528822a38363d001e Mon Sep 17 00:00:00 2001
From: Eno Compton <enocom@google.com>
Date: Mon, 31 Oct 2022 19:18:42 -0600
Subject: [PATCH] feat: improve refresh duration calculation

The change here updates the calculation of how long to wait before
starting a refresh cycle.

This is a port of https://github.com/GoogleCloudPlatform/alloydb-go-connector/pull/103.
---
 internal/cloudsql/instance.go      | 26 ++++++++++++------
 internal/cloudsql/instance_test.go | 44 ++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/internal/cloudsql/instance.go b/internal/cloudsql/instance.go
index 230cb5b6..6b797506 100644
--- a/internal/cloudsql/instance.go
+++ b/internal/cloudsql/instance.go
@@ -11,7 +11,6 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-
 package cloudsql
 
 import (
@@ -28,11 +27,6 @@ import (
 	sqladmin "google.golang.org/api/sqladmin/v1beta4"
 )
 
-const (
-	// refreshBuffer is the amount of time before a result expires to start a new refresh attempt.
-	refreshBuffer = 5 * time.Minute
-)
-
 var (
 	// Instance connection name is the format <PROJECT>:<REGION>:<INSTANCE>
 	// Additionally, we have to support legacy "domain-scoped" projects (e.g. "google.com:PROJECT")
@@ -274,6 +268,22 @@ func (i *Instance) result(ctx context.Context) (*refreshOperation, error) {
 	return res, nil
 }
 
+// refreshDuration returns the duration to wait before starting the next
+// refresh. Usually that duration will be half of the time until certificate
+// expiration.
+func refreshDuration(now, certExpiry time.Time) time.Duration {
+	d := certExpiry.Sub(now)
+	if d < time.Hour {
+		// Something is wrong with the certificate, refresh now.
+		if d < 5*time.Minute {
+			return 0
+		}
+		// Otherwise, wait five minutes before starting the refresh cycle.
+		return 5 * time.Minute
+	}
+	return d / 2
+}
+
 // scheduleRefresh schedules a refresh operation to be triggered after a given duration. The returned refreshOperation
 // can be used to either Cancel or Wait for the operations result.
 func (i *Instance) scheduleRefresh(d time.Duration) *refreshOperation {
@@ -306,8 +316,8 @@ func (i *Instance) scheduleRefresh(d time.Duration) *refreshOperation {
 			return
 		default:
 		}
-		nextRefresh := i.cur.expiry.Add(-refreshBuffer)
-		i.next = i.scheduleRefresh(time.Until(nextRefresh))
+		t := refreshDuration(time.Now(), i.cur.expiry)
+		i.next = i.scheduleRefresh(t)
 	})
 	return res
 }
diff --git a/internal/cloudsql/instance_test.go b/internal/cloudsql/instance_test.go
index 59e92d41..d86913b5 100644
--- a/internal/cloudsql/instance_test.go
+++ b/internal/cloudsql/instance_test.go
@@ -259,3 +259,47 @@ func TestClose(t *testing.T) {
 		t.Fatalf("failed to retrieve connect info: %v", err)
 	}
 }
+
+func TestRefreshDuration(t *testing.T) {
+	now := time.Now()
+	tcs := []struct {
+		desc   string
+		expiry time.Time
+		want   time.Duration
+	}{
+		{
+			desc:   "when expiration is greater than 1 hour",
+			expiry: now.Add(4 * time.Hour),
+			want:   2 * time.Hour,
+		},
+		{
+			desc:   "when expiration is equal to 1 hour",
+			expiry: now.Add(time.Hour),
+			want:   30 * time.Minute,
+		},
+		{
+			desc:   "when expiration is less than 1 hour, but greater than 5 minutes",
+			expiry: now.Add(6 * time.Minute),
+			want:   5 * time.Minute,
+		},
+		{
+			desc:   "when expiration is less than 5 minutes",
+			expiry: now.Add(4 * time.Minute),
+			want:   0,
+		},
+		{
+			desc:   "when expiration is now",
+			expiry: now,
+			want:   0,
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+			got := refreshDuration(now, tc.expiry)
+			// round to the second to remove millisecond differences
+			if got.Round(time.Second) != tc.want {
+				t.Fatalf("time until refresh: want = %v, got = %v", tc.want, got)
+			}
+		})
+	}
+}