Merge branch 'issue202'

supysonic/api/media.py

@@ -157,25 +157,50 @@ def stream_media():
             except OSError:
                 raise ServerError("Error while running the transcoding process")
 
+            if estimateContentLength == "true":
+                estimate = dst_bitrate * 1000 * res.duration // 8
+            else:
+                estimate = None
+
             def transcode():
+                while True:
+                    data = proc.stdout.read(8192)
+                    if not data:
+                        break
+                    yield data
+
+            def kill_processes():
+                if dec_proc != None:
+                    dec_proc.kill()
+                proc.kill()
+
+            def handle_transcoding():
                 try:
-                    while True:
-                        data = proc.stdout.read(8192)
-                        if not data:
-                            break
+                    sent = 0
+                    for data in transcode():
+                        sent += len(data)
                         yield data
-                except BaseException:
+                except (Exception, SystemExit, KeyboardInterrupt):
                     # Make sure child processes are always killed
-                    if dec_proc != None:
-                        dec_proc.kill()
-                    proc.kill()
+                    kill_processes()
                     raise
+                except GeneratorExit:
+                    # Try to transcode/send more data if we're close to the end.
+                    # The calling code have to support this as yielding more data
+                    # after a GeneratorExit would normally raise a RuntimeError.
+                    # Hopefully this generator is only used by the cache which
+                    # handles this.
+                    if estimate and sent >= estimate * 0.95:
+                        yield from transcode()
+                    else:
+                        kill_processes()
+                        raise
                 finally:
                     if dec_proc != None:
                         dec_proc.wait()
                     proc.wait()
 
-            resp_content = cache.set_generated(cache_key, transcode)
+            resp_content = cache.set_generated(cache_key, handle_transcoding)
 
             logger.info(
                 "Transcoding track {0.id} for user {1.id}. Source: {2} at {0.bitrate}kbps. Dest: {3} at {4}kbps".format(
@@ -183,10 +208,8 @@ def stream_media():
                 )
             )
             response = Response(resp_content, mimetype=dst_mimetype)
-            if estimateContentLength == "true":
-                response.headers.add(
-                    "Content-Length", dst_bitrate * 1000 * res.duration // 8
-                )
+            if estimate is not None:
+                response.headers.add("Content-Length", estimate)
     else:
         response = send_file(res.path, mimetype=dst_mimetype, conditional=True)
 

supysonic/cache.py

@@ -159,7 +159,7 @@ class Cache(object):
                     self._make_space(size, key=key)
                 os.replace(f.name, self._filepath(key))
                 self._record_file(key, size)
-        except Exception:
+        except BaseException:
             f.close()
             with contextlib.suppress(OSError):
                 os.remove(f.name)
@@ -183,9 +183,21 @@ class Cache(object):
         ...     print(x)
         """
         with self.set_fileobj(key) as f:
-            for data in gen_function():
-                f.write(data)
-                yield data
+            gen = gen_function()
+            try:
+                for data in gen:
+                    f.write(data)
+                    yield data
+            except GeneratorExit:
+                # Try to stop the generator but check it still wants to yield data.
+                # If it does allow caching of this data without forwarding it
+                try:
+                    f.write(gen.throw(GeneratorExit))
+                    for data in gen:
+                        f.write(data)
+                except StopIteration:
+                    # We stopped just at the end of the generator
+                    pass
 
     def get(self, key):
         """Return the path to the file where the cached data is stored"""

tests/api/test_transcoding.py

@@ -10,6 +10,7 @@
 import unittest
 import sys
 
+from flask import current_app
 from pony.orm import db_session
 
 from supysonic.db import Folder, Track
@@ -22,7 +23,6 @@ from .apitestbase import ApiTestBase
 class TranscodingTestCase(ApiTestBase):
     def setUp(self):
         super(TranscodingTestCase, self).setUp()
-        self._patch_client()
 
         with db_session:
             folder = FolderManager.add("Folder", "tests/assets/folder")
@@ -52,8 +52,10 @@ class TranscodingTestCase(ApiTestBase):
     )
     def test_direct_transcode(self):
         rv = self._stream(maxBitRate=96, estimateContentLength="true")
-        self.assertIn("tests/assets/folder/silence.mp3", rv.data)
-        self.assertTrue(rv.data.endswith("96"))
+        self.assertIn(b"tests/assets/folder/silence.mp3", rv.data)
+        self.assertTrue(rv.data.endswith(b"96"))
+        self.assertIn("Content-Length", rv.headers)
+        self.assertEqual(rv.content_length, 48000)  # 4s at 96kbps
 
     @unittest.skipIf(
         sys.platform == "win32",
@@ -61,10 +63,69 @@ class TranscodingTestCase(ApiTestBase):
     )
     def test_decode_encode(self):
         rv = self._stream(format="cat")
-        self.assertEqual(rv.data, "Pushing out some mp3 data...")
+        self.assertEqual(rv.data, b"Pushing out some mp3 data...")
 
         rv = self._stream(format="md5")
-        self.assertTrue(rv.data.startswith("dbb16c0847e5d8c3b1867604828cb50b"))
+        self.assertTrue(rv.data.startswith(b"dbb16c0847e5d8c3b1867604828cb50b"))
+
+    @unittest.skipIf(
+        sys.platform == "win32",
+        "Can't test transcoding on Windows because of a lack of simple commandline tools",
+    )
+    def test_mostly_transcoded_cached(self):
+        # See https://github.com/spl0k/supysonic/issues/202
+
+        rv = self._stream(maxBitRate=96, estimateContentLength="true", format="rnd")
+
+        read = 0
+        it = iter(rv.response)
+        # Read up to the estimated length
+        while read < 48000:
+            read += len(next(it))
+        rv.response.close()
+        rv.close()
+
+        key = "{}-96.rnd".format(self.trackid)
+        with self.app_context():
+            self.assertTrue(current_app.transcode_cache.has(key))
+            self.assertEqual(current_app.transcode_cache.size, 52000)
+
+    @unittest.skipIf(
+        sys.platform == "win32",
+        "Can't test transcoding on Windows because of a lack of simple commandline tools",
+    )
+    def test_partly_transcoded_cached(self):
+        rv = self._stream(maxBitRate=96, estimateContentLength="true", format="rnd")
+
+        # read one check of data then close the connection
+        next(iter(rv.response))
+        rv.response.close()
+        rv.close()
+
+        key = "{}-96.rnd".format(self.trackid)
+        with self.app_context():
+            self.assertFalse(current_app.transcode_cache.has(key))
+            self.assertEqual(current_app.transcode_cache.size, 0)
+
+    @unittest.skipIf(
+        sys.platform == "win32",
+        "Can't test transcoding on Windows because of a lack of simple commandline tools",
+    )
+    def test_last_chunk_close_transcoded_cached(self):
+        rv = self._stream(maxBitRate=96, estimateContentLength="true", format="rnd")
+
+        read = 0
+        it = iter(rv.response)
+        # Read up to the last chunk of data but keep the generator "alive"
+        while read < 52000:
+            read += len(next(it))
+        rv.response.close()
+        rv.close()
+
+        key = "{}-96.rnd".format(self.trackid)
+        with self.app_context():
+            self.assertTrue(current_app.transcode_cache.has(key))
+            self.assertEqual(current_app.transcode_cache.size, 52000)
 
 
 if __name__ == "__main__":

tests/testbase.py

@@ -25,6 +25,7 @@ class TestConfig(DefaultConfig):
     MIMETYPES = {"mp3": "audio/mpeg", "weirdextension": "application/octet-stream"}
     TRANSCODING = {
         "transcoder_mp3_mp3": "echo -n %srcpath %outrate",
+        "transcoder_mp3_rnd": "dd if=/dev/urandom bs=1kB count=52 status=none",
         "decoder_mp3": "echo -n Pushing out some mp3 data...",
         "encoder_cat": "cat -",
         "encoder_md5": "md5sum",
@@ -100,6 +101,9 @@ class TestBase(unittest.TestCase):
         self.client.get = patch_method(self.client.get)
         self.client.post = patch_method(self.client.post)
 
+    def app_context(self, *args, **kwargs):
+        return self.__app.app_context(*args, **kwargs)
+
     def request_context(self, *args, **kwargs):
         return self.__app.test_request_context(*args, **kwargs)