Add next_correlation_id_recv + send helper methods to BrokerConnection

1 files changed, 8 insertions, 0 deletions

diff --git a/kafka/conn.py b/kafka/conn.py
index 4516ddc..9907cb1 100644
--- a/kafka/conn.py
+++ b/kafka/conn.py
@@ -105,6 +105,14 @@ class BrokerConnection(local):
             return None
         return response
 
+    def next_correlation_id_recv(self):
+        if len(self.in_flight_requests) == 0:
+            return None
+        return self.in_flight_requests[0][0]
+
+    def next_correlation_id_send(self):
+        return self.correlation_id + 1
+
     def __getnewargs__(self):
         return (self.host, self.port, self.timeout)