3 files changed, 80 insertions, 4 deletions

diff --git a/media/libstagefright/CallbackDataSource.cpp b/media/libstagefright/CallbackDataSource.cpp
index 2e0745f..41f0175 100644
--- a/media/libstagefright/CallbackDataSource.cpp
+++ b/media/libstagefright/CallbackDataSource.cpp
@@ -70,9 +70,10 @@ ssize_t CallbackDataSource::readAt(off64_t offset, void* data, size_t size) {
         if (numRead == 0) {
             return totalNumRead;
         }
-        // Sanity check.
-        CHECK((size_t)numRead <= numToRead && numRead >= 0 &&
-                (size_t)numRead <= bufferSize);
+        if ((size_t)numRead > numToRead) {
+            return ERROR_OUT_OF_RANGE;
+        }
+        CHECK(numRead >= 0 && (size_t)numRead <= bufferSize);
         memcpy(((uint8_t*)data) + totalNumRead, mMemory->pointer(), numRead);
         numLeft -= numRead;
         totalNumRead += numRead;
@@ -94,4 +95,49 @@ status_t CallbackDataSource::getSize(off64_t *size) {
     return OK;
 }
 
+TinyCacheSource::TinyCacheSource(const sp<DataSource>& source)
+    : mSource(source), mCachedOffset(0), mCachedSize(0) {
+}
+
+status_t TinyCacheSource::initCheck() const {
+    return mSource->initCheck();
+}
+
+ssize_t TinyCacheSource::readAt(off64_t offset, void* data, size_t size) {
+    if (size >= kCacheSize) {
+        return mSource->readAt(offset, data, size);
+    }
+
+    // Check if the cache satisfies the read.
+    if (offset >= mCachedOffset && offset + size <= mCachedOffset + mCachedSize) {
+        memcpy(data, &mCache[offset - mCachedOffset], size);
+        return size;
+    }
+
+    // Fill the cache and copy to the caller.
+    const ssize_t numRead = mSource->readAt(offset, mCache, kCacheSize);
+    if (numRead <= 0) {
+        return numRead;
+    }
+    if ((size_t)numRead > kCacheSize) {
+        return ERROR_OUT_OF_RANGE;
+    }
+
+    mCachedSize = numRead;
+    mCachedOffset = offset;
+    CHECK(mCachedSize <= kCacheSize && mCachedOffset >= 0);
+    const size_t numToReturn = std::min(size, (size_t)numRead);
+    memcpy(data, mCache, numToReturn);
+
+    return numToReturn;
+}
+
+status_t TinyCacheSource::getSize(off64_t *size) {
+    return mSource->getSize(size);
+}
+
+uint32_t TinyCacheSource::flags() {
+    return mSource->flags();
+}
+
 } // namespace android
diff --git a/media/libstagefright/DataSource.cpp b/media/libstagefright/DataSource.cpp
index 6a89154..75ef288 100644
--- a/media/libstagefright/DataSource.cpp
+++ b/media/libstagefright/DataSource.cpp
@@ -283,7 +283,7 @@ sp<DataSource> DataSource::CreateMediaHTTP(const sp<IMediaHTTPService> &httpServ
 }
 
 sp<DataSource> DataSource::CreateFromIDataSource(const sp<IDataSource> &source) {
-    return new CallbackDataSource(source);
+    return new TinyCacheSource(new CallbackDataSource(source));
 }
 
 String8 DataSource::getMIMEType() const {
diff --git a/media/libstagefright/include/CallbackDataSource.h b/media/libstagefright/include/CallbackDataSource.h
index 678eb2e..1a21dd3 100644
--- a/media/libstagefright/include/CallbackDataSource.h
+++ b/media/libstagefright/include/CallbackDataSource.h
@@ -44,6 +44,36 @@ private:
     DISALLOW_EVIL_CONSTRUCTORS(CallbackDataSource);
 };
 
+
+// A caching DataSource that wraps a CallbackDataSource. For reads smaller
+// than kCacheSize it will read up to kCacheSize ahead and cache it.
+// This reduces the number of binder round trips to the IDataSource and has a significant
+// impact on time taken for filetype sniffing and metadata extraction.
+class TinyCacheSource : public DataSource {
+public:
+    TinyCacheSource(const sp<DataSource>& source);
+
+    virtual status_t initCheck() const;
+    virtual ssize_t readAt(off64_t offset, void* data, size_t size);
+    virtual status_t getSize(off64_t* size);
+    virtual uint32_t flags();
+
+private:
+    // 2kb comes from experimenting with the time-to-first-frame from a MediaPlayer
+    // with an in-memory MediaDataSource source on a Nexus 5. Beyond 2kb there was
+    // no improvement.
+    enum {
+        kCacheSize = 2048,
+    };
+
+    sp<DataSource> mSource;
+    uint8_t mCache[kCacheSize];
+    off64_t mCachedOffset;
+    size_t mCachedSize;
+
+    DISALLOW_EVIL_CONSTRUCTORS(TinyCacheSource);
+};
+
 }; // namespace android
 
 #endif // ANDROID_CALLBACKDATASOURCE_H