Improved the speed of set_prioritize_first_last in torrent.py

set_prioritize_first_last is rewritten for better performance, and is now only called when necessary. It should now properly set the priority for the beginning and end of each file in the torrent.

3 files changed, 173 insertions, 22 deletions

diff --git a/deluge/core/torrent.py b/deluge/core/torrent.py
index ba6ec285a..c5fd81955 100644
--- a/deluge/core/torrent.py
+++ b/deluge/core/torrent.py
@@ -232,8 +232,13 @@ class Torrent(object):
             "max_upload_speed": self.set_max_upload_speed,
             "prioritize_first_last_pieces": self.set_prioritize_first_last,
             "sequential_download": self.set_sequential_download
-
         }
+
+        # set_prioritize_first_last is called by set_file_priorities,
+        # so remove if file_priorities is set in options.
+        if "file_priorities" in options:
+            del OPTIONS_FUNCS["prioritize_first_last_pieces"]
+
         for (key, value) in options.items():
             if OPTIONS_FUNCS.has_key(key):
                 OPTIONS_FUNCS[key](value)
@@ -296,26 +301,44 @@ class Torrent(object):
 
     def set_prioritize_first_last(self, prioritize):
         self.options["prioritize_first_last_pieces"] = prioritize
-        if self.handle.has_metadata():
-            if self.options["compact_allocation"]:
-                log.debug("Setting first/last priority with compact "
-                          "allocation does not work!")
-                return
-
-            paths = {}
-            ti = self.handle.get_torrent_info()
-            for n in range(ti.num_pieces()):
-                slices = ti.map_block(n, 0, ti.piece_size(n))
-                for slice in slices:
-                    if self.handle.file_priority(slice.file_index):
-                        paths.setdefault(slice.file_index, []).append(n)
-
-            priorities = self.handle.piece_priorities()
-            for pieces in paths.itervalues():
-                two_percent = int(0.02*len(pieces)) or 1
-                for piece in pieces[:two_percent]+pieces[-two_percent:]:
-                    priorities[piece] = 7 if prioritize else 1
-            self.handle.prioritize_pieces(priorities)
+        if not prioritize:
+            # If we are turning off this option, call set_file_priorities to
+            # reset all the piece priorities
+            self.set_file_priorities(self.options["file_priorities"])
+            return
+        if not self.handle.has_metadata():
+            return
+        if self.options["compact_allocation"]:
+            log.debug("Setting first/last priority with compact "
+                      "allocation does not work!")
+            return
+        # A list of priorities for each piece in the torrent
+        priorities = self.handle.piece_priorities()
+        prioritized_pieces = []
+        ti = self.handle.get_torrent_info()
+        for i in range(ti.num_files()):
+            f = ti.file_at(i)
+            two_percent_bytes = int(0.02 * f.size)
+            # Get the pieces for the byte offsets
+            first_start = ti.map_file(i, 0, 0).piece
+            first_end = ti.map_file(i, two_percent_bytes, 0).piece
+            last_start = ti.map_file(i, f.size - two_percent_bytes, 0).piece
+            last_end = ti.map_file(i, max(f.size - 1, 0), 0).piece
+
+            first_end += 1
+            last_end += 1
+            prioritized_pieces.append((first_start, first_end))
+            prioritized_pieces.append((last_start, last_end))
+
+            # Creating two lists with priorites for the first/last pieces
+            # of this file, and insert the priorities into the list
+            first_list = [7] * (first_end - first_start)
+            last_list = [7] * (last_end - last_start)
+            priorities[first_start:first_end] = first_list
+            priorities[last_start:last_end] = last_list
+        # Setting the priorites for all the pieces of this torrent
+        self.handle.prioritize_pieces(priorities)
+        return prioritized_pieces, priorities
 
     def set_sequential_download(self, set_sequencial):
         self.options["sequential_download"] = set_sequencial
@@ -372,7 +395,8 @@ class Torrent(object):
             log.warning("File priorities were not set for this torrent")
 
         # Set the first/last priorities if needed
-        self.set_prioritize_first_last(self.options["prioritize_first_last_pieces"])
+        if self.options["prioritize_first_last_pieces"]:
+            self.set_prioritize_first_last(self.options["prioritize_first_last_pieces"])
 
     def set_trackers(self, trackers):
         """Sets trackers"""
diff --git a/deluge/tests/dir_with_6_files.torrent b/deluge/tests/dir_with_6_files.torrent
new file mode 100644
index 000000000..2c6b5fb02
--- /dev/null
+++ b/deluge/tests/dir_with_6_files.torrent
diff --git a/deluge/tests/test_torrent.py b/deluge/tests/test_torrent.py
new file mode 100644
index 000000000..f55b4d4e2
--- /dev/null
+++ b/deluge/tests/test_torrent.py
@@ -0,0 +1,127 @@
+from twisted.trial import unittest
+import os
+
+import deluge.core.torrent
+import test_torrent
+import common
+from deluge.core.rpcserver import RPCServer
+from deluge.core.core import Core
+
+deluge.core.torrent.component = test_torrent
+
+from deluge._libtorrent import lt
+
+import deluge.component as component
+from deluge.core.torrent import Torrent
+
+config_setup = False
+core = None
+rpcserver = None
+
+# This is called by torrent.py when calling component.get("...")
+def get(key):
+    if key is "Core":
+        return core
+    elif key is "RPCServer":
+        return rpcserver
+    else:
+        return None
+
+class TorrentTestCase(unittest.TestCase):
+
+    def setup_config(self):
+        global config_setup
+        config_setup = True
+        config_dir = common.set_tmp_config_dir()
+        core_config = deluge.config.Config("core.conf", defaults=deluge.core.preferencesmanager.DEFAULT_PREFS, config_dir=config_dir)
+        core_config.save()
+
+    def setUp(self):
+        self.setup_config()
+        global rpcserver
+        global core
+        rpcserver = RPCServer(listen=False)
+        core = Core()
+        self.session = lt.session()
+        self.torrent = None
+        return component.start()
+
+    def tearDown(self):
+        if self.torrent:
+            self.torrent.prev_status_cleanup_loop.stop()
+
+        def on_shutdown(result):
+            component._ComponentRegistry.components = {}
+        return component.shutdown().addCallback(on_shutdown)
+
+    def print_priority_list(self, priorities):
+        tmp = ''
+        for i, p in enumerate(priorities):
+            if i % 100 == 0:
+                print tmp
+                tmp = ''
+            tmp += "%s" % p
+        print tmp
+
+    def get_torrent_atp(self, filename):
+        filename = os.path.join(os.path.dirname(__file__), filename)
+        e = lt.bdecode(open(filename, 'rb').read())
+        info = lt.torrent_info(e)
+        atp = {"ti": info}
+        atp["save_path"] = os.getcwd()
+        atp["storage_mode"] = lt.storage_mode_t.storage_mode_sparse
+        atp["add_paused"] = False
+        atp["auto_managed"] = True
+        atp["duplicate_is_error"] = True
+        return atp
+
+    def test_set_prioritize_first_last(self):
+        piece_indexes = [(0,1), (0,1), (0,1), (0,1), (0,2), (50,52),
+                         (51,53), (110,112), (111,114), (200,203),
+                         (202,203), (212,213), (212,218), (457,463)]
+        self.run_test_set_prioritize_first_last("dir_with_6_files.torrent", piece_indexes)
+
+    def run_test_set_prioritize_first_last(self, torrent_file, prioritized_piece_indexes):
+        atp = self.get_torrent_atp(torrent_file)
+        handle = self.session.add_torrent(atp)
+
+        self.torrent = Torrent(handle, {})
+        priorities_original = handle.piece_priorities()
+        prioritized_pieces, new_priorites = self.torrent.set_prioritize_first_last(True)
+        priorities = handle.piece_priorities()
+        non_prioritized_pieces = list(range(len(priorities)))
+
+        # The prioritized indexes are the same as we expect
+        self.assertEquals(prioritized_pieces, prioritized_piece_indexes)
+
+        # Test the priority of the prioritized pieces
+        for first, last in prioritized_pieces:
+            for i in range(first, last):
+                if i in non_prioritized_pieces:
+                    non_prioritized_pieces.remove(i)
+                self.assertEquals(priorities[i], 7)
+
+        # Test the priority of all the non-prioritized pieces
+        for i in non_prioritized_pieces:
+            self.assertEquals(priorities[i], 1)
+
+        # The length of the list of new priorites is the same as the original
+        self.assertEquals(len(priorities_original), len(new_priorites))
+
+        #self.print_priority_list(priorities)
+
+    def test_set_prioritize_first_last_false(self):
+        atp = self.get_torrent_atp("dir_with_6_files.torrent")
+        handle = self.session.add_torrent(atp)
+        self.torrent = Torrent(handle, {})
+        # First set some pieces prioritized
+        self.torrent.set_prioritize_first_last(True)
+        # Reset pirorities
+        self.torrent.set_prioritize_first_last(False)
+        priorities = handle.piece_priorities()
+
+        # Test the priority of the prioritized pieces
+        for i in priorities:
+            self.assertEquals(priorities[i], 1)
+
+        #self.print_priority_list(priorities)