Review comments

1 files changed, 5 insertions, 12 deletions

diff --git a/synapse/rest/client/v2_alpha/sendtodevice.py b/synapse/rest/client/v2_alpha/sendtodevice.py
index 31167ba535..2ce038c6cd 100644
--- a/synapse/rest/client/v2_alpha/sendtodevice.py
+++ b/synapse/rest/client/v2_alpha/sendtodevice.py
@@ -19,8 +19,7 @@ from twisted.internet import defer
 
 from synapse.http import servlet
 from synapse.http.servlet import parse_json_object_from_request
-from synapse.rest.client.v1.transactions import HttpTransactionCache
-from synapse.util.async import ObservableDeferred
+from synapse.rest.client.transactions import HttpTransactionCache
 
 from ._base import client_v2_patterns
 
@@ -46,16 +45,10 @@ class SendToDeviceRestServlet(servlet.RestServlet):
 
     @defer.inlineCallbacks
     def on_PUT(self, request, message_type, txn_id):
-        try:
-            res_deferred = self.txns.get_client_transaction(request, txn_id)
-            res = yield res_deferred.observe()
-            defer.returnValue(res)
-        except KeyError:
-            pass
-
-        res_deferred = ObservableDeferred(self._put(request, message_type, txn_id))
-        self.txns.store_client_transaction(request, txn_id, res_deferred)
-        res = yield res_deferred.observe()
+        observable = self.txns.fetch_or_execute_request(
+            request, self._put, request, message_type, txn_id
+        )
+        res = yield observable.observe()
         defer.returnValue(res)
 
     @defer.inlineCallbacks