From eb06ae26609cbc46fa65e50c080508d53ec0b9c2 Mon Sep 17 00:00:00 2001
From: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>
Date: Fri, 7 May 2021 15:55:57 -0600
Subject: [PATCH] feat: force early refresh of instance info if connect fails
 (#19)

---
 dialer.go                         |  4 ++++
 internal/cloudsql/instance.go     | 12 ++++++++++++
 internal/cloudsql/refresh_test.go |  2 +-
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/dialer.go b/dialer.go
index 974b450b..ada614f8 100644
--- a/dialer.go
+++ b/dialer.go
@@ -140,6 +140,8 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 
 	conn, err := proxy.Dial(ctx, "tcp", addr)
 	if err != nil {
+		// refresh the instance info in case it caused the connection failure
+		i.ForceRefresh()
 		return nil, err
 	}
 	if c, ok := conn.(*net.TCPConn); ok {
@@ -152,6 +154,8 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 	}
 	tlsConn := tls.Client(conn, tlsCfg)
 	if err := tlsConn.Handshake(); err != nil {
+		// refresh the instance info in case it caused the handshake failure
+		i.ForceRefresh()
 		tlsConn.Close()
 		return nil, fmt.Errorf("handshake failed: %w", err)
 	}
diff --git a/internal/cloudsql/instance.go b/internal/cloudsql/instance.go
index 496cce1a..1b7d919e 100644
--- a/internal/cloudsql/instance.go
+++ b/internal/cloudsql/instance.go
@@ -147,6 +147,18 @@ func (i *Instance) ConnectInfo(ctx context.Context) (map[string]string, *tls.Con
 	return res.md.ipAddrs, res.tlsCfg, nil
 }
 
+// ForceRefresh triggers an immediate refresh operation to be scheduled and used for future connection attempts.
+func (i *Instance) ForceRefresh() {
+	i.resultGuard.Lock()
+	defer i.resultGuard.Unlock()
+	// If the next refresh hasn't started yet, we can cancel it and start an immediate one
+	if i.next.Cancel() {
+		i.next = i.scheduleRefresh(0)
+	}
+	// block all sequential connection attempts on the next refresh result
+	i.cur = i.next
+}
+
 // scheduleRefresh schedules a refresh operation to be triggered after a given duration. The returned refreshResult
 // can be used to either Cancel or Wait for the operations result.
 func (i *Instance) scheduleRefresh(d time.Duration) *refreshResult {
diff --git a/internal/cloudsql/refresh_test.go b/internal/cloudsql/refresh_test.go
index 4bff36a8..cb0378c7 100644
--- a/internal/cloudsql/refresh_test.go
+++ b/internal/cloudsql/refresh_test.go
@@ -40,7 +40,7 @@ func TestFetchMetadata(t *testing.T) {
 		t.Fatalf("%s", err)
 	}
 }
-func TestFetchEmpheralCert(t *testing.T) {
+func TestFetchEphemeralCert(t *testing.T) {
 	ctx := context.Background()
 
 	client, err := sqladmin.NewService(ctx)