Correctly ignore all incoming messages after the Terminate message

Once Terminate (`b'X'`) has been sent, there is no point in processing
any stray messages in the buffer, so discard them explicitly.

Author: elprans

asyncpg/exceptions/_base.py

@@ -12,7 +12,7 @@
 
 __all__ = ('PostgresError', 'FatalPostgresError', 'UnknownPostgresError',
            'InterfaceError', 'InterfaceWarning', 'PostgresLogMessage',
-           'InternalClientError', 'OutdatedSchemaCacheError')
+           'InternalClientError', 'OutdatedSchemaCacheError', 'ProtocolError')
 
 
 def _is_asyncpg_class(cls):

asyncpg/protocol/coreproto.pxd

@@ -17,6 +17,7 @@ cdef enum ProtocolState:
     PROTOCOL_FAILED = 1
     PROTOCOL_ERROR_CONSUME = 2
     PROTOCOL_CANCELLED = 3
+    PROTOCOL_TERMINATING = 4
 
     PROTOCOL_AUTH = 10
     PROTOCOL_PREPARE = 11

asyncpg/protocol/coreproto.pyx

@@ -118,9 +118,16 @@ cdef class CoreProtocol:
                     else:
                         self.buffer.discard_message()
 
+                elif state == PROTOCOL_TERMINATING:
+                    # The connection is being terminated.
+                    # discard all messages until connection
+                    # termination.
+                    self.buffer.discard_message()
+
                 else:
                     raise apg_exc.InternalClientError(
-                        'protocol is in an unknown state {}'.format(state))
+                        f'cannot process message {chr(mtype)!r}: '
+                        f'protocol is in an unexpected state {state!r}.')
 
             except Exception as ex:
                 self.result_type = RESULT_FAILED
@@ -637,8 +644,7 @@ cdef class CoreProtocol:
                     'cannot switch to "idle" state; '
                     'protocol is in the "failed" state')
             elif self.state == PROTOCOL_IDLE:
-                raise apg_exc.InternalClientError(
-                    'protocol is already in the "idle" state')
+                pass
             else:
                 self.state = new_state
 
@@ -648,6 +654,9 @@ cdef class CoreProtocol:
         elif new_state == PROTOCOL_CANCELLED:
             self.state = PROTOCOL_CANCELLED
 
+        elif new_state == PROTOCOL_TERMINATING:
+            self.state = PROTOCOL_TERMINATING
+
         else:
             if self.state == PROTOCOL_IDLE:
                 self.state = new_state
@@ -903,6 +912,7 @@ cdef class CoreProtocol:
     cdef _terminate(self):
         cdef WriteBuffer buf
         self._ensure_connected()
+        self._set_state(PROTOCOL_TERMINATING)
         buf = WriteBuffer.new_message(b'X')
         buf.end_message()
         self._write(buf)

asyncpg/protocol/protocol.pyx

@@ -805,6 +805,11 @@ cdef class BaseProtocol(CoreProtocol):
             elif self.state == PROTOCOL_COPY_IN_DATA:
                 self._on_result__copy_in(waiter)
 
+            elif self.state == PROTOCOL_TERMINATING:
+                # We are waiting for the connection to drop, so
+                # ignore any stray results at this point.
+                pass
+
             else:
                 raise apg_exc.InternalClientError(
                     'got result for unknown protocol state {}'.