POC connection hook

Author: pfreixes

internal/pool/pool.go

@@ -27,6 +27,23 @@ var timers = sync.Pool{
 	},
 }
 
+type ConnectEvent struct {
+        Err error
+}
+
+type PoolHook interface {
+        BeforeConnect(ctx context.Context) (context.Context, error)
+        AfterConnect(ctx context.Context, event ConnectEvent) error
+}
+
+type hooks struct  {
+	hooks []PoolHook
+}
+
+func (hs *hooks) AddHook(hook PoolHook) {
+	hs.hooks = append(hs.hooks, hook)
+}
+
 // Stats contains pool state information and accumulated stats.
 type Stats struct {
 	Hits     uint32 // number of times free connection was found in the pool
@@ -39,6 +56,8 @@ type Stats struct {
 }
 
 type Pooler interface {
+	AddHook(hook PoolHook)
+
 	NewConn(context.Context) (*Conn, error)
 	CloseConn(*Conn) error
 
@@ -70,6 +89,8 @@ type lastDialErrorWrap struct {
 }
 
 type ConnPool struct {
+	hooks
+
 	opt *Options
 
 	dialErrorsNum uint32 // atomic

internal/pool/pool_single.go

@@ -3,6 +3,7 @@ package pool
 import "context"
 
 type SingleConnPool struct {
+	hooks
 	pool      Pooler
 	cn        *Conn
 	stickyErr error

internal/pool/pool_sticky.go

@@ -34,6 +34,7 @@ func (e BadConnError) Unwrap() error {
 //------------------------------------------------------------------------------
 
 type StickyConnPool struct {
+	hooks
 	pool   Pooler
 	shared int32 // atomic
 

redis.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
 	"github.com/go-redis/redis/v8/internal"
 	"github.com/go-redis/redis/v8/internal/pool"
@@ -29,6 +30,11 @@ type Hook interface {
 	AfterProcessPipeline(ctx context.Context, cmds []Cmder) error
 }
 
+type fullHook interface {
+	Hook
+	pool.PoolHook
+}
+
 type hooks struct {
 	hooks []Hook
 }
@@ -45,6 +51,10 @@ func (hs hooks) clone() hooks {
 
 func (hs *hooks) AddHook(hook Hook) {
 	hs.hooks = append(hs.hooks, hook)
+	if hook, ok := hook.(fullHook); ok {
+		client := *(*Client)(unsafe.Pointer(hs))
+		client.baseClient.connPool.AddHook(hook)
+	}
 }
 
 func (hs hooks) process(