fix: update instance object rather than replace

kurtisvg · kurtisvg · commit ae0119363c09 · 2022-05-27T20:19:43.000Z
diff --git a/dialer.go b/dialer.go
@@ -186,7 +186,7 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 
 	var endInfo trace.EndSpanFunc
 	ctx, endInfo = trace.StartSpan(ctx, "cloud.google.com/go/cloudsqlconn/internal.InstanceInfo")
-	i, err := d.instance(instance, cfg.refreshCfg)
+	i, err := d.instance(instance, &cfg.refreshCfg)
 	if err != nil {
 		endInfo(err)
 		return nil, err
@@ -240,7 +240,7 @@ func (d *Dialer) Dial(ctx context.Context, instance string, opts ...DialOption)
 // corespond to one of the following types for the instance:
 // https://cloud.google.com/sql/docs/mysql/admin-api/rest/v1beta4/SqlDatabaseVersion
 func (d *Dialer) EngineVersion(ctx context.Context, instance string) (string, error) {
-	i, err := d.instance(instance, d.defaultDialCfg.refreshCfg)
+	i, err := d.instance(instance, nil)
 	if err != nil {
 		return "", err
 	}
@@ -258,7 +258,7 @@ func (d *Dialer) Warmup(ctx context.Context, instance string, opts ...DialOption
 	for _, opt := range opts {
 		opt(&cfg)
 	}
-	_, err := d.instance(instance, d.defaultDialCfg.refreshCfg)
+	_, err := d.instance(instance, &d.defaultDialCfg.refreshCfg)
 	return err
 }
 
@@ -301,30 +301,33 @@ func (d *Dialer) Close() error {
 	return nil
 }
 
-func (d *Dialer) instance(connName string, rCfg cloudsql.RefreshCfg) (*cloudsql.Instance, error) {
+// instance is a helper function for returning the appropriate instance object in a threadsafe way.
+// It will create a new instance object, modify the existing one, or leave it unchanged as needed.
+func (d *Dialer) instance(connName string, rCfg *cloudsql.RefreshCfg) (*cloudsql.Instance, error) {
 	// Check instance cache
 	d.lock.RLock()
 	i, ok := d.instances[connName]
 	d.lock.RUnlock()
 	// Check if the instance exists and that the refresh cfg is the same
-	if !ok || rCfg != i.RefreshCfg {
+	if !ok || (rCfg != nil && *rCfg != i.RefreshCfg) {
 		d.lock.Lock()
 		// Recheck to ensure instance wasn't created or changed between locks
 		i, ok = d.instances[connName]
-		if !ok || rCfg != i.RefreshCfg {
+		if !ok {
 			// Create a new instance
+			if rCfg == nil {
+				rCfg = &d.defaultDialCfg.refreshCfg
+			}
 			var err error
-			newI, err := cloudsql.NewInstance(connName, d.sqladmin, d.key, d.refreshTimeout, d.iamTokenSource, d.dialerID, rCfg)
+			i, err = cloudsql.NewInstance(connName, d.sqladmin, d.key, d.refreshTimeout, d.iamTokenSource, d.dialerID, *rCfg)
 			if err != nil {
 				d.lock.Unlock()
 				return nil, err
 			}
-			// If we created a new instance to match RefreshCfg, close the old one
-			if ok {
-				defer i.Close()
-			}
-			d.instances[connName] = newI
-			i = newI
+			d.instances[connName] = i
+		} else if rCfg != nil && *rCfg != i.RefreshCfg {
+			// Update the instance with the new refresh cfg
+			i.UpdateRefresh(*rCfg)
 		}
 		d.lock.Unlock()
 	}
diff --git a/internal/cloudsql/instance.go b/internal/cloudsql/instance.go
@@ -128,11 +128,11 @@ type Instance struct {
 	OpenConns uint64
 
 	connName
-	key        *rsa.PrivateKey
-	r          refresher
-	RefreshCfg RefreshCfg
+	key *rsa.PrivateKey
 
 	resultGuard sync.RWMutex
+	r           refresher
+	RefreshCfg  RefreshCfg
 	// cur represents the current refreshOperation that will be used to create connections. If a valid complete
 	// refreshOperation isn't available it's possible for cur to be equal to next.
 	cur *refreshOperation
@@ -161,10 +161,6 @@ func NewInstance(
 		return nil, err
 	}
 	ctx, cancel := context.WithCancel(context.Background())
-	var iamTs oauth2.TokenSource
-	if rCfg.UseIAMAuthN {
-		iamTs = ts
-	}
 	i := &Instance{
 		connName: cn,
 		key:      key,
@@ -173,7 +169,7 @@ func NewInstance(
 			30*time.Second,
 			2,
 			client,
-			iamTs,
+			ts,
 			dialerID,
 		),
 		ctx:    ctx,
@@ -224,6 +220,19 @@ func (i *Instance) InstanceEngineVersion(ctx context.Context) (string, error) {
 	return res.md.version, nil
 }
 
+func (i *Instance) UpdateRefresh(cfg RefreshCfg) {
+	i.resultGuard.Lock()
+	// Cancel any pending refreshes
+	i.cur.Cancel()
+	i.next.Cancel()
+	// update the refreshcfg as needed
+	i.RefreshCfg = cfg
+	// reschedule a new refresh immiediately
+	i.cur = i.scheduleRefresh(0)
+	i.next = i.cur
+	i.resultGuard.Unlock()
+}
+
 // ForceRefresh triggers an immediate refresh operation to be scheduled and used for future connection attempts.
 func (i *Instance) ForceRefresh() {
 	i.resultGuard.Lock()
@@ -254,7 +263,7 @@ func (i *Instance) scheduleRefresh(d time.Duration) *refreshOperation {
 	res := &refreshOperation{}
 	res.ready = make(chan struct{})
 	res.timer = time.AfterFunc(d, func() {
-		res.md, res.tlsCfg, res.expiry, res.err = i.r.performRefresh(i.ctx, i.connName, i.key)
+		res.md, res.tlsCfg, res.expiry, res.err = i.r.performRefresh(i.ctx, i.connName, i.key, i.RefreshCfg.UseIAMAuthN)
 		close(res.ready)
 
 		// Once the refresh is complete, update "current" with working result and schedule a new refresh
diff --git a/internal/cloudsql/refresh.go b/internal/cloudsql/refresh.go
@@ -270,7 +270,7 @@ func newRefresher(
 		dialerID:      dialerID,
 		clientLimiter: rate.NewLimiter(rate.Every(interval), burst),
 		client:        svc,
-		iamTS:         ts,
+		ts:            ts,
 	}
 }
 
@@ -286,12 +286,12 @@ type refresher struct {
 	clientLimiter *rate.Limiter
 	client        *sqladmin.Service
 
-	// iamTS is the TokenSource used for IAM DB AuthN. It is only set if IAM DB AuthN is being used.
-	iamTS oauth2.TokenSource
+	// ts is the TokenSource used for IAM DB AuthN.
+	ts oauth2.TokenSource
 }
 
 // performRefresh immediately performs a full refresh operation using the Cloud SQL Admin API.
-func (r refresher) performRefresh(ctx context.Context, cn connName, k *rsa.PrivateKey) (md metadata, c *tls.Config, expiry time.Time, err error) {
+func (r refresher) performRefresh(ctx context.Context, cn connName, k *rsa.PrivateKey, iamAuthN bool) (md metadata, c *tls.Config, expiry time.Time, err error) {
 	var refreshEnd trace.EndSpanFunc
 	ctx, refreshEnd = trace.StartSpan(ctx, "cloud.google.com/go/cloudsqlconn/internal.RefreshConnection",
 		trace.AddInstanceName(cn.String()),
@@ -337,7 +337,11 @@ func (r refresher) performRefresh(ctx context.Context, cn connName, k *rsa.Priva
 	ecC := make(chan ecRes, 1)
 	go func() {
 		defer close(ecC)
-		ec, err := fetchEphemeralCert(ctx, r.client, cn, k, r.iamTS)
+		var iamTS oauth2.TokenSource
+		if iamAuthN {
+			iamTS = r.ts
+		}
+		ec, err := fetchEphemeralCert(ctx, r.client, cn, k, iamTS)
 		ecC <- ecRes{ec, err}
 	}()
 
diff --git a/internal/cloudsql/refresh_test.go b/internal/cloudsql/refresh_test.go
@@ -60,7 +60,7 @@ func TestRefresh(t *testing.T) {
 	}()
 
 	r := newRefresher(time.Hour, 30*time.Second, 2, client, nil, "")
-	md, tlsCfg, gotExpiry, err := r.performRefresh(context.Background(), cn, RSAKey)
+	md, tlsCfg, gotExpiry, err := r.performRefresh(context.Background(), cn, RSAKey, false)
 	if err != nil {
 		t.Fatalf("PerformRefresh unexpectedly failed with error: %v", err)
 	}
@@ -101,22 +101,22 @@ func TestRefreshFailsFast(t *testing.T) {
 	defer cleanup()
 
 	r := newRefresher(time.Hour, 30*time.Second, 1, client, nil, "")
-	_, _, _, err = r.performRefresh(context.Background(), cn, RSAKey)
+	_, _, _, err = r.performRefresh(context.Background(), cn, RSAKey, false)
 	if err != nil {
 		t.Fatalf("expected no error, got = %v", err)
 	}
 
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
 	// context is canceled
-	_, _, _, err = r.performRefresh(ctx, cn, RSAKey)
+	_, _, _, err = r.performRefresh(ctx, cn, RSAKey, false)
 	if !errors.Is(err, context.Canceled) {
 		t.Fatalf("expected context.Canceled error, got = %v", err)
 	}
 
 	// force the rate limiter to throttle with a timed out context
 	ctx, _ = context.WithTimeout(context.Background(), time.Millisecond)
-	_, _, _, err = r.performRefresh(ctx, cn, RSAKey)
+	_, _, _, err = r.performRefresh(ctx, cn, RSAKey, false)
 
 	var wantErr *errtype.DialError
 	if !errors.As(err, &wantErr) {
@@ -186,7 +186,7 @@ func TestRefreshAdjustsCertExpiry(t *testing.T) {
 		t.Run(tc.desc, func(t *testing.T) {
 			ts := &fakeTokenSource{responses: tc.resps}
 			r := newRefresher(time.Hour, 30*time.Second, 1, client, ts, "")
-			_, _, gotExpiry, err := r.performRefresh(context.Background(), cn, RSAKey)
+			_, _, gotExpiry, err := r.performRefresh(context.Background(), cn, RSAKey, true)
 			if err != nil {
 				t.Fatalf("want no error, got = %v", err)
 			}
@@ -232,7 +232,7 @@ func TestRefreshWithIAMAuthErrors(t *testing.T) {
 		t.Run(tc.desc, func(t *testing.T) {
 			ts := &fakeTokenSource{responses: tc.resps}
 			r := newRefresher(time.Hour, 30*time.Second, 1, client, ts, "")
-			_, _, _, err = r.performRefresh(context.Background(), cn, RSAKey)
+			_, _, _, err = r.performRefresh(context.Background(), cn, RSAKey, true)
 			if err == nil {
 				t.Fatalf("expected get failed error, got = %v", err)
 			}
@@ -326,7 +326,7 @@ func TestRefreshWithFailedMetadataCall(t *testing.T) {
 			defer cleanup()
 
 			r := newRefresher(time.Hour, 30*time.Second, 1, client, nil, "")
-			_, _, _, err = r.performRefresh(context.Background(), cn, RSAKey)
+			_, _, _, err = r.performRefresh(context.Background(), cn, RSAKey, false)
 
 			if !errors.As(err, &tc.wantErr) {
 				t.Errorf("[%v] PerformRefresh failed with unexpected error, want = %T, got = %v", i, tc.wantErr, err)
@@ -392,7 +392,7 @@ func TestRefreshWithFailedEphemeralCertCall(t *testing.T) {
 		defer cleanup()
 
 		r := newRefresher(time.Hour, 30*time.Second, 1, client, nil, "")
-		_, _, _, err = r.performRefresh(context.Background(), cn, RSAKey)
+		_, _, _, err = r.performRefresh(context.Background(), cn, RSAKey, false)
 
 		if !errors.As(err, &tc.wantErr) {
 			t.Errorf("[%v] PerformRefresh failed with unexpected error, want = %T, got = %v", i, tc.wantErr, err)
@@ -419,7 +419,7 @@ func TestRefreshBuildsTLSConfig(t *testing.T) {
 	defer cleanup()
 
 	r := newRefresher(time.Hour, 30*time.Second, 1, client, nil, "")
-	_, tlsCfg, _, err := r.performRefresh(context.Background(), cn, RSAKey)
+	_, tlsCfg, _, err := r.performRefresh(context.Background(), cn, RSAKey, false)
 	if err != nil {
 		t.Fatalf("expected no error, got = %v", err)
 	}
