conncache - skip caching errors

cli/cmd/runtime/start.go

@@ -201,6 +201,7 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 			// Init runtime
 			opts := &runtime.Options{
 				ConnectionCacheSize:          conf.ConnectionCacheSize,
+				ConnectionCacheErrorTTL:      time.Minute,
 				MetastoreConnector:           "metastore",
 				QueryCacheSizeBytes:          conf.QueryCacheSizeBytes,
 				SecurityEngineCacheSize:      conf.SecurityEngineCacheSize,

cli/pkg/local/app.go

@@ -167,6 +167,7 @@ func NewApp(ctx context.Context, opts *AppOptions) (*App, error) {
 		SecurityEngineCacheSize:      1000,
 		ControllerLogBufferCapacity:  10000,
 		ControllerLogBufferSizeBytes: int64(datasize.MB * 16),
+		ConnectionCacheErrorTTL:      time.Minute,
 	}
 	rt, err := runtime.New(ctx, rtOpts, logger, opts.Ch.Telemetry(ctx), email.New(sender))
 	if err != nil {

runtime/connection_cache.go

@@ -40,6 +40,7 @@ func (r *Runtime) newConnectionCache() conncache.Cache {
 		MaxIdleConnections:   r.opts.ConnectionCacheSize,
 		OpenTimeout:          10 * time.Minute,
 		CloseTimeout:         10 * time.Minute,
+		ErrorTTL:             r.opts.ConnectionCacheErrorTTL,
 		CheckHangingInterval: time.Minute,
 		OpenFunc: func(ctx context.Context, cfg any) (conncache.Connection, error) {
 			x := cfg.(cachedConnectionConfig)

runtime/drivers/druid/druidsqldriver/druid_api_sql_driver.go

@@ -49,7 +49,10 @@ type sqlConnection struct {
 	dsn    string
 }
 
-var _ driver.QueryerContext = &sqlConnection{}
+var (
+	_ driver.QueryerContext = &sqlConnection{}
+	_ driver.Pinger         = &sqlConnection{}
+)
 
 func (c *sqlConnection) Prepare(query string) (driver.Stmt, error) {
 	return &stmt{
@@ -257,6 +260,15 @@ func createTransformer(columnType string) func(any) (any, error) {
 	}
 }
 
+func (c *sqlConnection) Ping(ctx context.Context) error {
+	rows, err := c.QueryContext(ctx, "SELECT 1", nil)
+	if err != nil {
+		return err
+	}
+	rows.Close()
+	return nil
+}
+
 type druidRows struct {
 	closer        io.ReadCloser
 	dec           *json.Decoder

runtime/pkg/conncache/conncache.go

@@ -62,6 +62,8 @@ type Options struct {
 	HangingFunc func(cfg any, open bool)
 	// Metrics are optional instruments for observability.
 	Metrics Metrics
+	// The maximum time the error stays in cache
+	ErrorTTL time.Duration
 }
 
 // Metrics are optional instruments for observability. If an instrument is nil, it will not be collected.
@@ -149,6 +151,8 @@ func (c *cacheImpl) Acquire(ctx context.Context, cfg any) (Connection, ReleaseFu
 	}
 
 	e, ok := c.entries[k]
+	reopen := ok && e.status == entryStatusOpen && e.err != nil && time.Since(e.since) > c.opts.ErrorTTL
+
 	if !ok {
 		e = &entry{cfg: cfg, since: time.Now()}
 		c.entries[k] = e
@@ -159,7 +163,7 @@ func (c *cacheImpl) Acquire(ctx context.Context, cfg any) (Connection, ReleaseFu
 
 	c.retainEntry(k, e)
 
-	if e.status == entryStatusOpen {
+	if e.status == entryStatusOpen && !reopen {
 		defer c.mu.Unlock()
 		if e.err != nil {
 			c.releaseEntry(k, e)
@@ -207,37 +211,37 @@ func (c *cacheImpl) Acquire(ctx context.Context, cfg any) (Connection, ReleaseFu
 		ch = make(chan struct{})
 		c.singleflight[k] = ch
 
-		e.status = entryStatusOpening
-		e.since = time.Now()
-		e.handle = nil
-		e.err = nil
+		if !reopen {
+			e.status = entryStatusOpening
+			e.since = time.Now()
+			e.handle = nil
+			e.err = nil
+		}
 
 		go func() {
-			start := time.Now()
-			var handle Connection
-			var err error
-			if c.opts.OpenTimeout == 0 {
-				handle, err = c.opts.OpenFunc(c.ctx, cfg)
-			} else {
-				ctx, cancel := context.WithTimeout(c.ctx, c.opts.OpenTimeout)
-				handle, err = c.opts.OpenFunc(ctx, cfg)
-				cancel()
-			}
-
-			if c.opts.Metrics.Opens != nil {
-				c.opts.Metrics.Opens.Add(c.ctx, 1)
+			var closeErr, err error
+			handle := e.handle
+			if reopen {
+				closeErr = c.closeConnection(e)
 			}
-			if c.opts.Metrics.OpenLatencyMS != nil {
-				c.opts.Metrics.OpenLatencyMS.Record(c.ctx, time.Since(start).Milliseconds())
+			if closeErr != nil || !reopen {
+				handle, err = c.openConnection(cfg)
 			}
 
 			c.mu.Lock()
 			defer c.mu.Unlock()
 
-			e.status = entryStatusOpen
-			e.since = time.Now()
-			e.handle = handle
-			e.err = err
+			if closeErr != nil {
+				e.status = entryStatusClosed
+				e.since = time.Now()
+				e.handle = nil
+				e.err = closeErr
+			} else {
+				e.status = entryStatusOpen
+				e.since = time.Now()
+				e.handle = handle
+				e.err = err
+			}
 
 			delete(c.singleflight, k)
 			close(ch)
@@ -278,6 +282,28 @@ func (c *cacheImpl) Acquire(ctx context.Context, cfg any) (Connection, ReleaseFu
 	return e.handle, c.releaseFunc(k, e), nil
 }
 
+func (c *cacheImpl) openConnection(cfg any) (Connection, error) {
+	start := time.Now()
+	var handle Connection
+	var err error
+	if c.opts.OpenTimeout == 0 {
+		handle, err = c.opts.OpenFunc(c.ctx, cfg)
+	} else {
+		ctx, cancel := context.WithTimeout(c.ctx, c.opts.OpenTimeout)
+		handle, err = c.opts.OpenFunc(ctx, cfg)
+		cancel()
+	}
+
+	if c.opts.Metrics.Opens != nil {
+		c.opts.Metrics.Opens.Add(c.ctx, 1)
+	}
+	if c.opts.Metrics.OpenLatencyMS != nil {
+		c.opts.Metrics.OpenLatencyMS.Record(c.ctx, time.Since(start).Milliseconds())
+	}
+
+	return handle, err
+}
+
 func (c *cacheImpl) EvictWhere(predicate func(cfg any) bool) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -358,21 +384,7 @@ func (c *cacheImpl) beginClose(k string, e *entry) {
 	e.since = time.Now()
 
 	go func() {
-		start := time.Now()
-		var err error
-		if e.handle != nil {
-			err = e.handle.Close()
-		}
-		if err == nil {
-			err = errors.New("conncache: connection closed")
-		}
-
-		if c.opts.Metrics.Closes != nil {
-			c.opts.Metrics.Closes.Add(c.ctx, 1)
-		}
-		if c.opts.Metrics.CloseLatencyMS != nil {
-			c.opts.Metrics.CloseLatencyMS.Record(c.ctx, time.Since(start).Milliseconds())
-		}
+		err := c.closeConnection(e)
 
 		c.mu.Lock()
 		defer c.mu.Unlock()
@@ -389,6 +401,25 @@ func (c *cacheImpl) beginClose(k string, e *entry) {
 	}()
 }
 
+func (c *cacheImpl) closeConnection(e *entry) error {
+	start := time.Now()
+	var err error
+	if e.handle != nil {
+		err = e.handle.Close()
+	}
+	if err == nil {
+		err = errors.New("conncache: connection closed")
+	}
+
+	if c.opts.Metrics.Closes != nil {
+		c.opts.Metrics.Closes.Add(c.ctx, 1)
+	}
+	if c.opts.Metrics.CloseLatencyMS != nil {
+		c.opts.Metrics.CloseLatencyMS.Record(c.ctx, time.Since(start).Milliseconds())
+	}
+	return err
+}
+
 func (c *cacheImpl) lruEvictionHandler(key, value any) {
 	k := key.(string)
 	e := value.(*entry)

runtime/pkg/conncache/conncache_test.go

@@ -27,6 +27,96 @@ func (c *mockConn) Close() error {
 	return nil
 }
 
+func TestImpl_count_1(t *testing.T) {
+	opens := atomic.Int64{}
+	c := New(Options{
+		MaxIdleConnections: 2,
+		OpenFunc: func(ctx context.Context, cfg any) (Connection, error) {
+			opens.Add(1)
+			return &mockConn{cfg: cfg.(string)}, nil
+		},
+		KeyFunc: func(cfg any) string {
+			return cfg.(string)
+		},
+	})
+	ci := c.(*cacheImpl)
+	_, r1, err := c.Acquire(context.Background(), "foo")
+	require.NoError(t, err)
+	time.Sleep(time.Second)
+	require.Equal(t, 1, len(ci.entries))
+	require.Equal(t, 1, ci.entries["foo"].refs)
+	require.Equal(t, 0, ci.lru.Len())
+
+	_, r2, err := c.Acquire(context.Background(), "foo")
+	require.Equal(t, 1, len(ci.entries))
+	require.Equal(t, 0, ci.lru.Len())
+	r1()
+	require.Equal(t, 1, len(ci.entries))
+	require.Equal(t, 0, ci.lru.Len())
+	r2()
+	require.Equal(t, 1, len(ci.entries))
+	require.Equal(t, 1, ci.lru.Len())
+
+	require.NoError(t, c.Close(context.Background()))
+}
+
+func TestImpl_err(t *testing.T) {
+	opens := atomic.Int64{}
+	c := New(Options{
+		MaxIdleConnections: 1,
+		OpenFunc: func(ctx context.Context, cfg any) (Connection, error) {
+			opens.Add(1)
+			return nil, fmt.Errorf("err")
+		},
+		KeyFunc: func(cfg any) string {
+			return cfg.(string)
+		},
+		ErrorTTL: time.Minute,
+	})
+	ci := c.(*cacheImpl)
+	_, _, err := c.Acquire(context.Background(), "foo")
+	require.Error(t, err)
+	time.Sleep(time.Second)
+	require.Equal(t, 1, len(ci.entries))
+	require.Equal(t, 1, ci.lru.Len())
+
+	_, _, err = c.Acquire(context.Background(), "foo")
+	require.Error(t, err)
+	require.Equal(t, 1, len(ci.entries))
+	require.Equal(t, 1, ci.lru.Len())
+
+	require.NoError(t, c.Close(context.Background()))
+}
+
+func TestImpl_lru(t *testing.T) {
+	opens := atomic.Int64{}
+	c := New(Options{
+		MaxIdleConnections: 1,
+		OpenFunc: func(ctx context.Context, cfg any) (Connection, error) {
+			opens.Add(1)
+			return &mockConn{cfg: cfg.(string)}, nil
+		},
+		KeyFunc: func(cfg any) string {
+			return cfg.(string)
+		},
+	})
+	ci := c.(*cacheImpl)
+	_, r1, err := c.Acquire(context.Background(), "foo")
+	require.NoError(t, err)
+	_, r2, err := c.Acquire(context.Background(), "foo2")
+	require.Equal(t, 2, len(ci.entries))
+	require.Equal(t, 0, ci.lru.Len())
+	r1()
+	require.Equal(t, 2, len(ci.entries))
+	require.Equal(t, 1, ci.lru.Len())
+	r2()
+	time.Sleep(time.Second)
+	require.Equal(t, 1, len(ci.entries))
+	require.Equal(t, 1, ci.lru.Len())
+
+	require.NoError(t, c.Close(context.Background()))
+}
+
 func TestBasic(t *testing.T) {
 	opens := atomic.Int64{}
 

runtime/runtime.go

@@ -29,6 +29,7 @@ type Options struct {
 	ControllerLogBufferSizeBytes int64
 	AllowHostAccess              bool
 	DataDir                      string
+	ConnectionCacheErrorTTL      time.Duration
 }
 
 type Runtime struct {

runtime/testruntime/testruntime.go

@@ -9,6 +9,7 @@ import (
 	goruntime "runtime"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/c2h5oh/datasize"
 	"github.com/joho/godotenv"
@@ -60,6 +61,7 @@ func New(t TestingT) *runtime.Runtime {
 			},
 		},
 		ConnectionCacheSize:          100,
+		ConnectionCacheErrorTTL:      time.Minute,
 		QueryCacheSizeBytes:          int64(datasize.MB * 100),
 		SecurityEngineCacheSize:      100,
 		ControllerLogBufferCapacity:  10000,