prefetch flag docs

3 files changed, 99 insertions, 42 deletions

diff --git a/src/leap/mail/imap/mailbox.py b/src/leap/mail/imap/mailbox.py
index 84bfa54..f319bf0 100644
--- a/src/leap/mail/imap/mailbox.py
+++ b/src/leap/mail/imap/mailbox.py
@@ -90,6 +90,8 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
 
     next_uid_lock = threading.Lock()
 
+    _fdoc_primed = {}
+
     def __init__(self, mbox, soledad, memstore, rw=1):
         """
         SoledadMailbox constructor. Needs to get passed a name, plus a
@@ -129,6 +131,7 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
         if self._memstore:
             self.prime_known_uids_to_memstore()
             self.prime_last_uid_to_memstore()
+            self.prime_flag_docs_to_memstore()
 
     @property
     def listeners(self):
@@ -279,6 +282,16 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
         known_uids = self.messages.all_soledad_uid_iter()
         self._memstore.set_known_uids(self.mbox, known_uids)
 
+    def prime_flag_docs_to_memstore(self):
+        """
+        Prime memstore with all the flags documents.
+        """
+        primed = self._fdoc_primed.get(self.mbox, False)
+        if not primed:
+            all_flag_docs = self.messages.get_all_soledad_flag_docs()
+            self._memstore.load_flag_docs(self.mbox, all_flag_docs)
+            self._fdoc_primed[self.mbox] = True
+
     def getUIDValidity(self):
         """
         Return the unique validity identifier for this mailbox.
@@ -606,7 +619,7 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
         messages_asked = self._bound_seq(messages_asked)
         seq_messg = self._filter_msg_seq(messages_asked)
 
-        all_flags = self.messages.all_flags()
+        all_flags = self._memstore.all_flags(self.mbox)
         result = ((msgid, flagsPart(
             msgid, all_flags.get(msgid, tuple()))) for msgid in seq_messg)
         return result
@@ -833,7 +846,7 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
                 self._memstore.create_message(
                     self.mbox, uid_next,
                     MessageWrapper(
-                        new_fdoc, hdoc.content),
+                        new_fdoc, hdoc),
                     observer=observer,
                     notify_on_disk=False)
 
@@ -860,6 +873,7 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
             logger.warning("Tried to copy a MSG with no fdoc")
             return
         new_fdoc = copy.deepcopy(fdoc.content)
+        copy_hdoc = copy.deepcopy(hdoc.content)
         fdoc_chash = new_fdoc[fields.CONTENT_HASH_KEY]
 
         # XXX is this hitting the db??? --- probably.
@@ -867,7 +881,7 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
         dest_fdoc = memstore.get_fdoc_from_chash(
             fdoc_chash, self.mbox)
         exist = dest_fdoc and not empty(dest_fdoc.content)
-        return exist, new_fdoc, hdoc
+        return exist, new_fdoc, copy_hdoc
 
     # convenience fun
 
diff --git a/src/leap/mail/imap/memorystore.py b/src/leap/mail/imap/memorystore.py
index 8deddda..00cf2cc 100644
--- a/src/leap/mail/imap/memorystore.py
+++ b/src/leap/mail/imap/memorystore.py
@@ -49,6 +49,11 @@ logger = logging.getLogger(__name__)
 # soledad storage, in seconds.
 SOLEDAD_WRITE_PERIOD = 10
 
+FDOC = MessagePartType.fdoc.key
+HDOC = MessagePartType.hdoc.key
+CDOCS = MessagePartType.cdocs.key
+DOCS_ID = MessagePartType.docs_id.key
+
 
 @contextlib.contextmanager
 def set_bool_flag(obj, att):
@@ -104,6 +109,11 @@ class MemoryStore(object):
         self._write_period = write_period
 
         # Internal Storage: messages
+        # TODO this probably will have better access times if we
+        # use msg_store[mbox][uid] insted of the current key scheme.
+        """
+        key is str(mbox,uid)
+        """
         self._msg_store = {}
 
         # Sizes
@@ -297,11 +307,6 @@ class MemoryStore(object):
         key = mbox, uid
         msg_dict = message.as_dict()
 
-        FDOC = MessagePartType.fdoc.key
-        HDOC = MessagePartType.hdoc.key
-        CDOCS = MessagePartType.cdocs.key
-        DOCS_ID = MessagePartType.docs_id.key
-
         try:
             store = self._msg_store[key]
         except KeyError:
@@ -580,6 +585,44 @@ class MemoryStore(object):
         if self._permanent_store:
             self._permanent_store.write_last_uid(mbox, value)
 
+    def load_flag_docs(self, mbox, flag_docs):
+        """
+        Load the flag documents for the given mbox.
+        Used during initial flag docs prefetch.
+
+        :param mbox: the mailbox
+        :type mbox: str or unicode
+        :param flag_docs: a dict with the content for the flag docs.
+        :type flag_docs: dict
+        """
+        # We can do direct assignments cause we know this will only
+        # be called during initialization of the mailbox.
+        msg_store = self._msg_store
+        for uid in flag_docs:
+            key = mbox, uid
+            msg_store[key] = {}
+            msg_store[key][FDOC] = ReferenciableDict(flag_docs[uid])
+
+    def all_flags(self, mbox):
+        """
+        Return a dictionary with all the flags for a given mbox.
+
+        :param mbox: the mailbox
+        :type mbox: str or unicode
+        :rtype: dict
+        """
+        flags_dict = {}
+        uids = self.get_uids(mbox)
+        store = self._msg_store
+        for uid in uids:
+            key = mbox, uid
+            try:
+                flags = store[key][FDOC][fields.FLAGS_KEY]
+                flags_dict[uid] = flags
+            except KeyError:
+                continue
+        return flags_dict
+
     # Counting sheeps...
 
     def count_new_mbox(self, mbox):
diff --git a/src/leap/mail/imap/messages.py b/src/leap/mail/imap/messages.py
index 89beaaa..3ba9d1b 100644
--- a/src/leap/mail/imap/messages.py
+++ b/src/leap/mail/imap/messages.py
@@ -919,7 +919,10 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
         if existing_uid:
             uid = existing_uid
             msg = self.get_msg_by_uid(uid)
-            reactor.callLater(0, msg.setFlags, (fields.DELETED_FLAG,), -1)
+
+            # TODO this cannot be deferred, this has to block.
+            #reactor.callLater(0, msg.setFlags, (fields.DELETED_FLAG,), -1)
+            msg.setFlags((fields.DELETED_FLAG,), -1)
             reactor.callLater(0, observer.callback, uid)
             return
 
@@ -1221,49 +1224,46 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
         """
         Return an iterator through the UIDs of all messages, from memory.
         """
-        if self.memstore is not None:
-            mem_uids = self.memstore.get_uids(self.mbox)
-            soledad_known_uids = self.memstore.get_soledad_known_uids(
-                self.mbox)
-            combined = tuple(set(mem_uids).union(soledad_known_uids))
-            return combined
+        mem_uids = self.memstore.get_uids(self.mbox)
+        soledad_known_uids = self.memstore.get_soledad_known_uids(
+            self.mbox)
+        combined = tuple(set(mem_uids).union(soledad_known_uids))
+        return combined
 
-    # XXX MOVE to memstore
-    def all_flags(self):
+    def get_all_soledad_flag_docs(self):
         """
-        Return a dict with all flags documents for this mailbox.
-        """
-        # XXX get all from memstore and cache it there
-        # FIXME should get all uids, get them fro memstore,
-        # and get only the missing ones from disk.
+        Return a dict with the content of all the flag documents
+        in soledad store for the given mbox.
 
+        :param mbox: the mailbox
+        :type mbox: str or unicode
+        :rtype: dict
+        """
+        # XXX we really could return a reduced version with
+        # just {'uid': (flags-tuple,) since the prefetch is
+        # only oriented to get the flag tuples.
         all_flags = dict(((
             doc.content[self.UID_KEY],
-            doc.content[self.FLAGS_KEY]) for doc in
+            dict(doc.content)) for doc in
             self._soledad.get_from_index(
                 fields.TYPE_MBOX_IDX,
                 fields.TYPE_FLAGS_VAL, self.mbox)))
-        if self.memstore is not None:
-            uids = self.memstore.get_uids(self.mbox)
-            docs = ((uid, self.memstore.get_message(self.mbox, uid))
-                    for uid in uids)
-            for uid, doc in docs:
-                all_flags[uid] = doc.fdoc.content[self.FLAGS_KEY]
-
         return all_flags
 
-    def all_flags_chash(self):
-        """
-        Return a dict with the content-hash for all flag documents
-        for this mailbox.
-        """
-        all_flags_chash = dict(((
-            doc.content[self.UID_KEY],
-            doc.content[self.CONTENT_HASH_KEY]) for doc in
-            self._soledad.get_from_index(
-                fields.TYPE_MBOX_IDX,
-                fields.TYPE_FLAGS_VAL, self.mbox)))
-        return all_flags_chash
+    # XXX Move to memstore too. But we don't need it really, since
+    # we can cache the headers docs too.
+    #def all_flags_chash(self):
+        #"""
+        #Return a dict with the content-hash for all flag documents
+        #for this mailbox.
+        #"""
+        #all_flags_chash = dict(((
+            #doc.content[self.UID_KEY],
+            #doc.content[self.CONTENT_HASH_KEY]) for doc in
+            #self._soledad.get_from_index(
+                #fields.TYPE_MBOX_IDX,
+                #fields.TYPE_FLAGS_VAL, self.mbox)))
+        #return all_flags_chash
 
     def all_headers(self):
         """