Properly cleanup connections closed by remote

elprans · elprans · commit 4393a159e661 · 2018-11-09T11:24:14.000-05:00
When a connection is terminated by the remote peer, asyncpg must not forget to perform all the necessary client-side cleanup procedures. Fixes: MagicStack#385.
diff --git a/asyncpg/protocol/protocol.pyx b/asyncpg/protocol/protocol.pyx
@@ -111,7 +111,10 @@ cdef class BaseProtocol(CoreProtocol):
 self.conref = weakref.ref(connection)
 
 cdef get_connection(self):
- return self.conref()
+ if self.conref is not None:
+ return self.conref()
+ else:
+ return None
 
 def get_server_pid(self):
 return self.backend_pid
@@ -867,6 +870,11 @@ cdef class BaseProtocol(CoreProtocol):
 # terminated or due to another error;
 # Throw an error in any awaiting waiter.
 self.closing = True
+ # Cleanup the connection resources, including, possibly,
+ # releasing the pool holder.
+ con = self.get_connection()
+ if con is not None:
+ con._cleanup()
 self._handle_waiter_on_connection_lost(exc)
 
 cdef _write(self, buf):
diff --git a/tests/test_pool.py b/tests/test_pool.py
@@ -881,6 +881,31 @@ async def test_pool_init_and_use_race(self):
 await pool_task
 await pool.close()
 
+ async def test_pool_remote_close(self):
+ pool = await self.create_pool(min_size=1, max_size=1)
+ backend_pid_fut = self.loop.create_future()
+
+ async def worker():
+ async with pool.acquire() as conn:
+ pool_backend_pid = await conn.fetchval(
+ 'SELECT pg_backend_pid()')
+ backend_pid_fut.set_result(pool_backend_pid)
+ await asyncio.sleep(0.2, loop=self.loop)
+
+ task = self.loop.create_task(worker())
+ try:
+ conn = await self.connect()
+ backend_pid = await backend_pid_fut
+ await conn.execute('SELECT pg_terminate_backend($1)', backend_pid)
+ finally:
+ await conn.close()
+
+ await task
+
+ # Check that connection_lost has released the pool holder.
+ conn = await pool.acquire(timeout=0.1)
+ await pool.release(conn)
+
 
 @unittest.skipIf(os.environ.get('PGHOST'), 'using remote cluster for testing')
 class TestHotStandby(tb.ClusterTestCase):

-Original file line number
+Diff line change
 self.conref = weakref.ref(connection)
  cdef get_connection(self):
 -returnself.conref()
 +ifself.conref isnotNone:
 +returnself.conref()
 +else:
 +returnNone
 defget_server_pid(self):
 returnself.backend_pid
 # terminated or due to another error;
 # Throw an error in any awaiting waiter.
 self.closing =True
 +# Cleanup the connection resources, including, possibly,
 +# releasing the pool holder.
 + con =self.get_connection()
 +if con isnotNone:
 + con._cleanup()
 self._handle_waiter_on_connection_lost(exc)
  cdef _write(self, buf):
-Original file line number
+Diff line change
 awaitpool_task
 awaitpool.close()
 +asyncdeftest_pool_remote_close(self):
 +pool=awaitself.create_pool(min_size=1, max_size=1)
 +backend_pid_fut=self.loop.create_future()
++
 +asyncdefworker():
 +asyncwithpool.acquire() asconn:
 +pool_backend_pid=awaitconn.fetchval(
 +'SELECT pg_backend_pid()')
 +backend_pid_fut.set_result(pool_backend_pid)
 +awaitasyncio.sleep(0.2, loop=self.loop)
++
 +task=self.loop.create_task(worker())
 +try:
 +conn=awaitself.connect()
 +backend_pid=awaitbackend_pid_fut
 +awaitconn.execute('SELECT pg_terminate_backend($1)', backend_pid)
 +finally:
 +awaitconn.close()
++
 +awaittask
++
 +# Check that connection_lost has released the pool holder.
 +conn=awaitpool.acquire(timeout=0.1)
 +awaitpool.release(conn)
++
 @unittest.skipIf(os.environ.get('PGHOST'), 'using remote cluster for testing')
 classTestHotStandby(tb.ClusterTestCase):