OnPut is called when a connection is returned to the pool
(ctx context.Context, conn *pool.Conn)
| 138 | |
| 139 | // OnPut is called when a connection is returned to the pool |
| 140 | func (ph *PoolHook) OnPut(ctx context.Context, conn *pool.Conn) (shouldPool bool, shouldRemove bool, err error) { |
| 141 | // first check if we should handoff for faster rejection |
| 142 | if !conn.ShouldHandoff() { |
| 143 | // Default behavior (no handoff): pool the connection |
| 144 | return true, false, nil |
| 145 | } |
| 146 | |
| 147 | // check pending handoff to not queue the same connection twice |
| 148 | if ph.workerManager.isHandoffPending(conn) { |
| 149 | // Default behavior (pending handoff): pool the connection |
| 150 | return true, false, nil |
| 151 | } |
| 152 | |
| 153 | if err := ph.workerManager.queueHandoff(conn); err != nil { |
| 154 | // Failed to queue handoff, remove the connection |
| 155 | internal.Logger.Printf(ctx, logs.FailedToQueueHandoff(conn.GetID(), err)) |
| 156 | // Don't pool, remove connection, no error to caller |
| 157 | return false, true, nil |
| 158 | } |
| 159 | |
| 160 | // Check if handoff was already processed by a worker before we can mark it as queued |
| 161 | if !conn.ShouldHandoff() { |
| 162 | // Handoff was already processed - this is normal and the connection should be pooled |
| 163 | return true, false, nil |
| 164 | } |
| 165 | |
| 166 | if err := conn.MarkQueuedForHandoff(); err != nil { |
| 167 | // If marking fails, check if handoff was processed in the meantime |
| 168 | if !conn.ShouldHandoff() { |
| 169 | // Handoff was processed - this is normal, pool the connection |
| 170 | return true, false, nil |
| 171 | } |
| 172 | // Other error - remove the connection |
| 173 | return false, true, nil |
| 174 | } |
| 175 | internal.Logger.Printf(ctx, logs.MarkedForHandoff(conn.GetID())) |
| 176 | return true, false, nil |
| 177 | } |
| 178 | |
| 179 | func (ph *PoolHook) OnRemove(_ context.Context, conn *pool.Conn, _ error) { |
| 180 | if tracker, ok := ph.operationsManager.(maintNotificationsConnTracker); ok && conn != nil { |