From 1a9510bb84d79f6ff78d32390195bc97ed9a439e Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Mon, 17 Aug 2015 10:40:23 +0100
Subject: [PATCH] Implement a batched presence_handler.get_state and use it

---
 synapse/handlers/message.py  | 18 ++++-------
 synapse/handlers/presence.py | 63 ++++++++++++++++++++++++++++++++++++
 synapse/storage/presence.py  |  6 ++--
 3 files changed, 73 insertions(+), 14 deletions(-)

diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index 29e81085d1..f12465fa2c 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -460,20 +460,14 @@ class MessageHandler(BaseHandler):
 
         @defer.inlineCallbacks
         def get_presence():
-            presence_defs = yield defer.DeferredList(
-                [
-                    presence_handler.get_state(
-                        target_user=UserID.from_string(m.user_id),
-                        auth_user=auth_user,
-                        as_event=True,
-                        check_auth=False,
-                    )
-                    for m in room_members
-                ],
-                consumeErrors=True,
+            states = yield presence_handler.get_states(
+                target_users=[UserID.from_string(m.user_id) for m in room_members],
+                auth_user=auth_user,
+                as_event=True,
+                check_auth=False,
             )
 
-            defer.returnValue([p for success, p in presence_defs if success])
+            defer.returnValue(states.values())
 
         receipts_handler = self.hs.get_handlers().receipts_handler
 
diff --git a/synapse/handlers/presence.py b/synapse/handlers/presence.py
index 341a516da2..33d76efe08 100644
--- a/synapse/handlers/presence.py
+++ b/synapse/handlers/presence.py
@@ -232,6 +232,69 @@ class PresenceHandler(BaseHandler):
         else:
             defer.returnValue(state)
 
+    @defer.inlineCallbacks
+    def get_states(self, target_users, auth_user, as_event=False, check_auth=True):
+        try:
+            local_users, remote_users = partitionbool(
+                target_users,
+                lambda u: self.hs.is_mine(u)
+            )
+
+            if check_auth:
+                for u in local_users:
+                    visible = yield self.is_presence_visible(
+                        observer_user=auth_user,
+                        observed_user=u
+                    )
+
+                    if not visible:
+                        raise SynapseError(404, "Presence information not visible")
+
+            results = {}
+            if local_users:
+                for u in local_users:
+                    if u in self._user_cachemap:
+                        results[u] = self._user_cachemap[u].get_state()
+
+                local_to_user = {u.localpart: u for u in local_users}
+
+                states = yield self.store.get_presence_states(
+                    [u.localpart for u in local_users if u not in results]
+                )
+
+                for local_part, state in states.items():
+                    res = {"presence": state["state"]}
+                    if "status_msg" in state and state["status_msg"]:
+                        res["status_msg"] = state["status_msg"]
+                    results[local_to_user[local_part]] = res
+
+            for u in remote_users:
+                # TODO(paul): Have remote server send us permissions set
+                results[u] = self._get_or_offline_usercache(u).get_state()
+
+            for state in results.values():
+                if "last_active" in state:
+                    state["last_active_ago"] = int(
+                        self.clock.time_msec() - state.pop("last_active")
+                    )
+
+            if as_event:
+                for user, state in results.items():
+                    content = state
+                    content["user_id"] = user.to_string()
+
+                    if "last_active" in content:
+                        content["last_active_ago"] = int(
+                            self._clock.time_msec() - content.pop("last_active")
+                        )
+
+                    results[user] = {"type": "m.presence", "content": content}
+        except:
+            logger.exception(":(")
+            raise
+
+        defer.returnValue(results)
+
     @defer.inlineCallbacks
     @log_function
     def set_state(self, target_user, auth_user, state):
diff --git a/synapse/storage/presence.py b/synapse/storage/presence.py
index f351b76a70..15d98198e2 100644
--- a/synapse/storage/presence.py
+++ b/synapse/storage/presence.py
@@ -50,13 +50,15 @@ class PresenceStore(SQLBaseStore):
         def f(txn):
             results = {}
             for user_localpart in user_localparts:
-                results[user_localpart] = self._simple_select_one_txn(
+                res = self._simple_select_one_txn(
                     txn,
                     table="presence",
                     keyvalues={"user_id": user_localpart},
                     retcols=["state", "status_msg", "mtime"],
-                    desc="get_presence_state",
+                    allow_none=True,
                 )
+                if res:
+                    results[user_localpart] = res
 
             return results