chore: Refactor dialer cache concurrency logic. Part of #842.

Author: hessjcg

dialer.go

@@ -122,12 +122,41 @@ type monitoredCache struct {
 	connectionInfoCache
 }
 
+func (c *monitoredCache) cleanupClosed() {
+	// Remove closed sockets from cache.openSockets
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	var newOpenSockets []*instrumentedConn
+	for _, s := range c.openSockets {
+		if !s.closed {
+			newOpenSockets = append(newOpenSockets, s)
+		}
+	}
+	c.openSockets = newOpenSockets
+
+}
+func (c *monitoredCache) closeMonitored() {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	if c.openConns != nil {
+		for _, socket := range c.openSockets {
+			if !socket.closed {
+				socket.Close()
+			}
+		}
+		atomic.StoreUint64(c.openConns, 0)
+	}
+	if c.connectionInfoCache != nil {
+		c.connectionInfoCache.Close()
+	}
+}
+
 // A Dialer is used to create connections to Cloud SQL instances.
 //
 // Use NewDialer to initialize a Dialer.
 type Dialer struct {
-	lock           sync.RWMutex
-	cache          map[instance.ConnName]monitoredCache
+	cache          *DialerCache
 	keyGenerator   *keyGenerator
 	refreshTimeout time.Duration
 	// closed reports if the dialer has been closed.
@@ -273,7 +302,7 @@ func NewDialer(ctx context.Context, opts ...Option) (*Dialer, error) {
 
 	d := &Dialer{
 		closed:            make(chan struct{}),
-		cache:             make(map[instance.ConnName]monitoredCache),
+		cache:             newDialerCache(cfg.logger),
 		lazyRefresh:       cfg.lazyRefresh,
 		keyGenerator:      g,
 		refreshTimeout:    cfg.refreshTimeout,
@@ -431,16 +460,14 @@ func (d *Dialer) removeCached(
 	ctx context.Context,
 	i instance.ConnName, c connectionInfoCache, err error,
 ) {
+	d.cache.remove(i)
 	d.logger.Debugf(
 		ctx,
 		"[%v] Removing connection info from cache: %v",
 		i.String(),
 		err,
 	)
-	d.lock.Lock()
-	defer d.lock.Unlock()
 	c.Close()
-	delete(d.cache, i)
 }
 
 // validClientCert checks that the ephemeral client certificate retrieved from
@@ -476,6 +503,7 @@ func (d *Dialer) EngineVersion(ctx context.Context, icn string) (string, error)
 	if err != nil {
 		return "", err
 	}
+	// Create a connectionInfoCache without adding it to the cache.
 	c, err := d.connectionInfoCache(ctx, cn, &d.defaultDialConfig.useIAMAuthN)
 	if err != nil {
 		return "", err
@@ -574,94 +602,60 @@ func (d *Dialer) Close() error {
 	default:
 	}
 	close(d.closed)
-	d.lock.Lock()
-	defer d.lock.Unlock()
+
+	d.cache.replaceAll(func(cn instance.ConnName, c monitoredCache) (instance.ConnName, monitoredCache) {
+		c.closeMonitored()                           // close the monitoredCache
+		return instance.ConnName{}, monitoredCache{} // Remove from cache
+	})
+
 	if d.domainNameTicker != nil {
 		d.domainNameTicker.Stop()
 	}
-	for _, i := range d.cache {
-		i.Close()
-	}
 	return nil
 }
 
 func (d *Dialer) pollDomainNames(ctx context.Context) {
-	type cacheEntry struct {
-		cn    instance.ConnName
-		cache monitoredCache
-	}
-
-	// List all the cache entries created with a domain name
-	d.lock.RLock()
-	cacheEntries := make([]cacheEntry, 0, len(d.cache))
-	for cn, cache := range d.cache {
-
-		// Ignore cache entries that were not opened by domain name.
+	// Check if domain changed.
+	d.cache.replaceAll(func(cn instance.ConnName, c monitoredCache) (instance.ConnName, monitoredCache) {
+		// No domain set, do nothing.
 		if cn.DomainName() == "" {
-			continue
-		}
-
-		cacheEntries = append(cacheEntries, cacheEntry{cn: cn, cache: cache})
-	}
-	d.lock.RUnlock()
-
-	for _, entry := range cacheEntries {
-		// Remove closed sockets from cache.openSockets
-		entry.cache.lock.Lock()
-		var newOpenSockets []*instrumentedConn
-		for _, s := range entry.cache.openSockets {
-			if !s.closed {
-				newOpenSockets = append(newOpenSockets, s)
-			}
+			return cn, c // no change
 		}
-		entry.cache.openSockets = newOpenSockets
-		entry.cache.lock.Unlock()
 
 		// Resolve the domain name.
-		newCn, err := d.resolver.Resolve(ctx, entry.cn.DomainName())
+		newCn, err := d.resolver.Resolve(ctx, cn.DomainName())
 
-		// the domain name no longer resolves to a valid instance
+		// The domain name no longer resolves, remove from cache.
 		if err != nil {
-			d.logger.Debugf(ctx, "[failover] unable to resolve DNS for instance %s: %v", entry.cn.DomainName(), err)
+			d.logger.Debugf(ctx, "[failover] unable to resolve DNS for instance %s: %v", cn.DomainName(), err)
+			c.closeMonitored()          // Close the instance
+			return cn, monitoredCache{} // remove from cache
 		}
 
-		// The domain name points to a different instance.
-		if newCn != entry.cn {
-			d.logger.Debugf(ctx, "domain name %s changed from old instance %s to new instance %s",
-				entry.cn.DomainName(), entry.cn.String(), newCn.String())
+		// The domain name points to a different instance, replace.
+		if newCn != cn {
+			d.logger.Debugf(ctx, "[failover] domain name %s changed from old instance %s to new instance %s",
+				cn.DomainName(), cn.String(), newCn.String())
 
-			d.closeDomainNameChanged(ctx, entry.cn, entry.cache,
-				fmt.Errorf("domain name %s changed from old instance %s to new instance %s",
-					entry.cn.DomainName(), entry.cn.String(), newCn.String()))
-			// preload the new cache entry
-			b := entry.cache.UseIAMAuthN()
-			d.connectionInfoCache(ctx, newCn, &b)
-		}
-	}
-
-}
+			// Close the old instance
+			go c.closeMonitored()
 
-func (d *Dialer) closeDomainNameChanged(ctx context.Context, cn instance.ConnName, cache monitoredCache, err error) {
-	d.removeCached(ctx, cn, cache, err)
-	if atomic.LoadUint64(cache.openConns) > 0 {
-		for _, socket := range cache.openSockets {
-			if !socket.closed {
-				socket.Close()
+			// preload the new cache entry
+			b := c.UseIAMAuthN()
+			newCache, err := d.createConnectionInfoCache(ctx, newCn, &b)
+			if err != nil {
+				return cn, monitoredCache{} // create instance failed, remove from cache
 			}
+
+			// replace with newCn with updated domain and new connectionInfoCache
+			return newCn, newCache
 		}
-	}
-}
 
-func (d *Dialer) findByDomainName(dn string) (instance.ConnName, monitoredCache, bool) {
+		// Nothing changed, clean up any closed sockets from the list of sockets.
+		c.cleanupClosed()
+		return cn, c
+	})
 
-	d.lock.RLock()
-	defer d.lock.RUnlock()
-	for cn, cache := range d.cache {
-		if cn.DomainName() == dn {
-			return cn, cache, true
-		}
-	}
-	return instance.ConnName{}, monitoredCache{}, false
 }
 
 // connectionInfoCache is a helper function for returning the appropriate
@@ -670,60 +664,49 @@ func (d *Dialer) findByDomainName(dn string) (instance.ConnName, monitoredCache,
 func (d *Dialer) connectionInfoCache(
 	ctx context.Context, cn instance.ConnName, useIAMAuthN *bool,
 ) (monitoredCache, error) {
-	d.lock.RLock()
-	c, ok := d.cache[cn]
-	d.lock.RUnlock()
-
-	if !ok {
-		// Check if the domain name was previously associated with a different
-		// instance, and if so, close that cache.
-		if cn.DomainName() != "" {
-			oldCn, c, ok := d.findByDomainName(cn.DomainName())
-			if ok {
-				d.closeDomainNameChanged(ctx, oldCn, c, fmt.Errorf(
-					"domain name %s changed from old instance %s to new instance %s",
-					cn.DomainName(), oldCn.String(), cn.String()))
-			}
-		}
 
-		// Create a new connectionInfoCache
-		d.lock.Lock()
-		defer d.lock.Unlock()
-		// Recheck to ensure instance wasn't created or changed between locks
-		c, ok = d.cache[cn]
-		if !ok {
-			var useIAMAuthNDial bool
-			if useIAMAuthN != nil {
-				useIAMAuthNDial = *useIAMAuthN
-			}
-			d.logger.Debugf(ctx, "[%v] Connection info added to cache", cn.String())
-			k, err := d.keyGenerator.rsaKey()
-			if err != nil {
-				return monitoredCache{}, err
-			}
-			var cache connectionInfoCache
-			if d.lazyRefresh {
-				cache = cloudsql.NewLazyRefreshCache(
-					cn,
-					d.logger,
-					d.sqladmin, k,
-					d.refreshTimeout, d.iamTokenSource,
-					d.dialerID, useIAMAuthNDial,
-				)
-			} else {
-				cache = cloudsql.NewRefreshAheadCache(
-					cn,
-					d.logger,
-					d.sqladmin, k,
-					d.refreshTimeout, d.iamTokenSource,
-					d.dialerID, useIAMAuthNDial,
-				)
-			}
-			var count uint64
-			c = monitoredCache{openConns: &count, connectionInfoCache: cache}
-			d.cache[cn] = c
-		}
+	c, oldC, err := d.cache.get(cn, func() (monitoredCache, error) {
+		return d.createConnectionInfoCache(ctx, cn, useIAMAuthN)
+	})
+
+	oldC.closeMonitored()
+
+	return c, err
+}
+
+func (d *Dialer) createConnectionInfoCache(
+	ctx context.Context, cn instance.ConnName, useIAMAuthN *bool,
+) (monitoredCache, error) {
+
+	var useIAMAuthNDial bool
+	if useIAMAuthN != nil {
+		useIAMAuthNDial = *useIAMAuthN
+	}
+	d.logger.Debugf(ctx, "[%v] Connection info created", cn.String())
+	k, err := d.keyGenerator.rsaKey()
+	if err != nil {
+		return monitoredCache{}, err
+	}
+	var cache connectionInfoCache
+	if d.lazyRefresh {
+		cache = cloudsql.NewLazyRefreshCache(
+			cn,
+			d.logger,
+			d.sqladmin, k,
+			d.refreshTimeout, d.iamTokenSource,
+			d.dialerID, useIAMAuthNDial,
+		)
+	} else {
+		cache = cloudsql.NewRefreshAheadCache(
+			cn,
+			d.logger,
+			d.sqladmin, k,
+			d.refreshTimeout, d.iamTokenSource,
+			d.dialerID, useIAMAuthNDial,
+		)
 	}
+	var count uint64
+	c := monitoredCache{openConns: &count, connectionInfoCache: cache}
 
 	c.UpdateRefresh(useIAMAuthN)