From 4c79a63fd76e982e5e60b22c7efd15b6e3cf9915 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Wed, 23 Nov 2016 10:40:44 +0000
Subject: [PATCH] Explicit federation ack

---
 synapse/app/federation_sender.py |  5 ++++-
 synapse/federation/send_queue.py | 13 +++++++++++--
 synapse/replication/resource.py  | 15 ++++++++++-----
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/synapse/app/federation_sender.py b/synapse/app/federation_sender.py
index ba2b4c261..dcdbe79a1 100644
--- a/synapse/app/federation_sender.py
+++ b/synapse/app/federation_sender.py
@@ -235,7 +235,10 @@ class FederationSenderHandler(object):
     @defer.inlineCallbacks
     def stream_positions(self):
         stream_id = yield self.store.get_federation_out_pos("federation")
-        defer.returnValue({"federation": stream_id})
+        defer.returnValue({
+            "federation": stream_id,
+            "federation_ack": stream_id,
+        })
 
     @defer.inlineCallbacks
     def process_replication(self, result):
diff --git a/synapse/federation/send_queue.py b/synapse/federation/send_queue.py
index ed2b03fad..5c9f7a86f 100644
--- a/synapse/federation/send_queue.py
+++ b/synapse/federation/send_queue.py
@@ -213,7 +213,15 @@ class FederationRemoteSendQueue(object):
     def get_current_token(self):
         return self.pos - 1
 
-    def get_replication_rows(self, token, limit):
+    def get_replication_rows(self, token, limit, federation_ack=None):
+        """
+        Args:
+            token (int)
+            limit (int)
+            federation_ack (int): Optional. The position where the worker is
+                explicitly acknowledged it has handled. Allows us to drop
+                data from before that point
+        """
         # TODO: Handle limit.
 
         # To handle restarts where we wrap around
@@ -224,7 +232,8 @@ class FederationRemoteSendQueue(object):
 
         # There should be only one reader, so lets delete everything its
         # acknowledged its seen.
-        self._clear_queue_before_pos(token)
+        if federation_ack:
+            self._clear_queue_before_pos(federation_ack)
 
         # Fetch changed presence
         keys = self.presence_changed.keys()
diff --git a/synapse/replication/resource.py b/synapse/replication/resource.py
index b05ca6271..cb9697e37 100644
--- a/synapse/replication/resource.py
+++ b/synapse/replication/resource.py
@@ -171,8 +171,13 @@ class ReplicationResource(Resource):
         }
         request_streams["streams"] = parse_string(request, "streams")
 
+        federation_ack = parse_integer(request, "federation_ack", None)
+
         def replicate():
-            return self.replicate(request_streams, limit)
+            return self.replicate(
+                request_streams, limit,
+                federation_ack=federation_ack
+            )
 
         writer = yield self.notifier.wait_for_replication(replicate, timeout)
         result = writer.finish()
@@ -190,7 +195,7 @@ class ReplicationResource(Resource):
         finish_request(request)
 
     @defer.inlineCallbacks
-    def replicate(self, request_streams, limit):
+    def replicate(self, request_streams, limit, federation_ack=None):
         writer = _Writer()
         current_token = yield self.current_replication_token()
         logger.debug("Replicating up to %r", current_token)
@@ -209,7 +214,7 @@ class ReplicationResource(Resource):
         yield self.caches(writer, current_token, limit, request_streams)
         yield self.to_device(writer, current_token, limit, request_streams)
         yield self.public_rooms(writer, current_token, limit, request_streams)
-        self.federation(writer, current_token, limit, request_streams)
+        self.federation(writer, current_token, limit, request_streams, federation_ack)
         self.streams(writer, current_token, request_streams)
 
         logger.debug("Replicated %d rows", writer.total)
@@ -473,7 +478,7 @@ class ReplicationResource(Resource):
                 "position", "room_id", "visibility"
             ), position=upto_token)
 
-    def federation(self, writer, current_token, limit, request_streams):
+    def federation(self, writer, current_token, limit, request_streams, federation_ack):
         if self.config.send_federation:
             return
 
@@ -483,7 +488,7 @@ class ReplicationResource(Resource):
 
         if federation is not None and federation != current_position:
             federation_rows = self.federation_sender.get_replication_rows(
-                federation, limit,
+                federation, limit, federation_ack=federation_ack,
             )
             upto_token = _position_from_rows(federation_rows, current_position)
             writer.write_header_and_rows("federation", federation_rows, (