OnPut is called when a connection is returned to the pool. This hook checks if the connection needs re-authentication. If so, it schedules a background goroutine to perform the re-auth asynchronously. The goroutine: 1. Waits for a worker slot (semaphore) 2. Acquires the connection (waits until not
(_ context.Context, conn *pool.Conn)
| 137 | // |
| 138 | // Thread-safe: Called concurrently by multiple goroutines returning connections. |
| 139 | func (r *ReAuthPoolHook) OnPut(_ context.Context, conn *pool.Conn) (bool, bool, error) { |
| 140 | if conn == nil { |
| 141 | // noop |
| 142 | return true, false, nil |
| 143 | } |
| 144 | connID := conn.GetID() |
| 145 | // Check if reauth is needed and get the function with proper locking |
| 146 | r.shouldReAuthLock.RLock() |
| 147 | reAuthFn, ok := r.shouldReAuth[connID] |
| 148 | r.shouldReAuthLock.RUnlock() |
| 149 | |
| 150 | if ok { |
| 151 | // Acquire both locks to atomically move from shouldReAuth to scheduledReAuth |
| 152 | // This prevents race conditions where OnGet might miss the transition |
| 153 | r.shouldReAuthLock.Lock() |
| 154 | r.scheduledLock.Lock() |
| 155 | r.scheduledReAuth[connID] = true |
| 156 | delete(r.shouldReAuth, connID) |
| 157 | r.scheduledLock.Unlock() |
| 158 | r.shouldReAuthLock.Unlock() |
| 159 | go func() { |
| 160 | r.workers.AcquireBlocking() |
| 161 | // safety first |
| 162 | if conn == nil || (conn != nil && conn.IsClosed()) { |
| 163 | r.workers.Release() |
| 164 | return |
| 165 | } |
| 166 | defer func() { |
| 167 | if rec := recover(); rec != nil { |
| 168 | // once again - safety first |
| 169 | internal.Logger.Printf(context.Background(), "panic in reauth worker: %v", rec) |
| 170 | } |
| 171 | r.scheduledLock.Lock() |
| 172 | delete(r.scheduledReAuth, connID) |
| 173 | r.scheduledLock.Unlock() |
| 174 | r.workers.Release() |
| 175 | }() |
| 176 | |
| 177 | // Create timeout context for connection acquisition |
| 178 | // This prevents indefinite waiting if the connection is stuck |
| 179 | ctx, cancel := context.WithTimeout(context.Background(), r.reAuthTimeout) |
| 180 | defer cancel() |
| 181 | |
| 182 | // Try to acquire the connection for re-authentication |
| 183 | // We need to ensure the connection is IDLE (not IN_USE) before transitioning to UNUSABLE |
| 184 | // This prevents re-authentication from interfering with active commands |
| 185 | // Use AwaitAndTransition to wait for the connection to become IDLE |
| 186 | stateMachine := conn.GetStateMachine() |
| 187 | if stateMachine == nil { |
| 188 | // No state machine - should not happen, but handle gracefully |
| 189 | reAuthFn(pool.ErrConnUnusableTimeout) |
| 190 | return |
| 191 | } |
| 192 | |
| 193 | // Use predefined slice to avoid allocation |
| 194 | _, err := stateMachine.AwaitAndTransition(ctx, pool.ValidFromIdle(), pool.StateUnusable) |
| 195 | if err != nil { |
| 196 | // Timeout or other error occurred, cannot acquire connection |
nothing calls this directly
no test coverage detected