feat: Allow prefetch init segments (shaka-project#5825)

lib/media/segment_prefetch.js

@@ -40,8 +40,9 @@ shaka.media.SegmentPrefetch = class {
     this.fetchDispatcher_ = fetchDispatcher;
 
     /**
-     * @private {!Map.<shaka.media.SegmentReference,
-     *           !shaka.media.SegmentPrefetchOperation>}
+     * @private {!Map.<
+     *        !(shaka.media.SegmentReference|shaka.media.InitSegmentReference),
+     *        !shaka.media.SegmentPrefetchOperation>}
      */
     this.segmentPrefetchMap_ = new Map();
 
@@ -52,9 +53,10 @@ shaka.media.SegmentPrefetch = class {
    * Fetch next segments ahead of current segment.
    *
    * @param {(!shaka.media.SegmentReference)} startReference
+   * @param {boolean=} skipFirst
    * @public
    */
-  prefetchSegments(startReference) {
+  prefetchSegments(startReference, skipFirst = false) {
     goog.asserts.assert(this.prefetchLimit_ > 0,
         'SegmentPrefetch can not be used when prefetchLimit <= 0.');
 
@@ -70,6 +72,14 @@ shaka.media.SegmentPrefetch = class {
       return;
     }
     let reference = startReference;
+    if (skipFirst) {
+      reference = iterator.next().value;
+      if (reference &&
+          reference.startTime == startReference.startTime &&
+          reference.endTime == startReference.endTime) {
+        reference = null;
+      }
+    }
     while (this.segmentPrefetchMap_.size < this.prefetchLimit_ &&
             reference != null) {
       if (!this.segmentPrefetchMap_.has(reference)) {
@@ -83,18 +93,45 @@ shaka.media.SegmentPrefetch = class {
     }
   }
 
+  /**
+   * Fetch init segment.
+   *
+   * @param {!shaka.media.InitSegmentReference} initSegmentReference
+   * @public
+   */
+  prefetchInitSegment(initSegmentReference) {
+    goog.asserts.assert(this.prefetchLimit_ > 0,
+        'SegmentPrefetch can not be used when prefetchLimit <= 0.');
+
+    const logPrefix = shaka.media.SegmentPrefetch.logPrefix_(this.stream_);
+    if (!this.stream_.segmentIndex) {
+      shaka.log.info(logPrefix, 'missing segmentIndex');
+      return;
+    }
+
+    if (this.segmentPrefetchMap_.size < this.prefetchLimit_) {
+      if (!this.segmentPrefetchMap_.has(initSegmentReference)) {
+        const segmentPrefetchOperation =
+          new shaka.media.SegmentPrefetchOperation(this.fetchDispatcher_);
+        segmentPrefetchOperation.dispatchFetch(
+            initSegmentReference, this.stream_);
+        this.segmentPrefetchMap_.set(
+            initSegmentReference, segmentPrefetchOperation);
+      }
+    }
+  }
+
   /**
    * Get the result of prefetched segment if already exists.
-   * @param {(!shaka.media.SegmentReference)} reference
+   * @param {!(shaka.media.SegmentReference|shaka.media.InitSegmentReference)}
+   *        reference
    * @param {?function(BufferSource):!Promise=} streamDataCallback
    * @return {?shaka.net.NetworkingEngine.PendingRequest} op
    * @public
    */
   getPrefetchedSegment(reference, streamDataCallback) {
     goog.asserts.assert(this.prefetchLimit_ > 0,
         'SegmentPrefetch can not be used when prefetchLimit <= 0.');
-    goog.asserts.assert(reference instanceof shaka.media.SegmentReference,
-        'getPrefetchedSegment is only used for shaka.media.SegmentReference.');
 
     const logPrefix = shaka.media.SegmentPrefetch.logPrefix_(this.stream_);
 
@@ -104,16 +141,30 @@ shaka.media.SegmentPrefetch = class {
         segmentPrefetchOperation.setStreamDataCallback(streamDataCallback);
       }
       this.segmentPrefetchMap_.delete(reference);
-      shaka.log.debug(
-          logPrefix,
-          'reused prefetched segment at time:', reference.startTime,
-          'mapSize', this.segmentPrefetchMap_.size);
+      if (reference instanceof shaka.media.SegmentReference) {
+        shaka.log.debug(
+            logPrefix,
+            'reused prefetched segment at time:', reference.startTime,
+            'mapSize', this.segmentPrefetchMap_.size);
+      } else {
+        shaka.log.debug(
+            logPrefix,
+            'reused prefetched init segment at time, mapSize',
+            this.segmentPrefetchMap_.size);
+      }
       return segmentPrefetchOperation.getOperation();
     } else {
-      shaka.log.debug(
-          logPrefix,
-          'missed segment at time:', reference.startTime,
-          'mapSize', this.segmentPrefetchMap_.size);
+      if (reference instanceof shaka.media.SegmentReference) {
+        shaka.log.debug(
+            logPrefix,
+            'reused prefetched segment at time:', reference.startTime,
+            'mapSize', this.segmentPrefetchMap_.size);
+      } else {
+        shaka.log.debug(
+            logPrefix,
+            'reused prefetched init segment at time, mapSize',
+            this.segmentPrefetchMap_.size);
+      }
       return null;
     }
   }
@@ -176,7 +227,8 @@ shaka.media.SegmentPrefetch = class {
 
   /**
    * Remove a segment from prefetch map and abort it.
-   * @param {(!shaka.media.SegmentReference)} reference
+   * @param {!(shaka.media.SegmentReference|shaka.media.InitSegmentReference)}
+   *        reference
    * @private
    */
   abortPrefetchedSegment_(reference) {
@@ -185,9 +237,13 @@ shaka.media.SegmentPrefetch = class {
     this.segmentPrefetchMap_.delete(reference);
     if (segmentPrefetchOperation) {
       segmentPrefetchOperation.abort();
-      shaka.log.info(
-          logPrefix,
-          'pop and abort prefetched segment at time:', reference.startTime);
+      if (reference instanceof shaka.media.SegmentReference) {
+        shaka.log.info(
+            logPrefix,
+            'pop and abort prefetched segment at time:', reference.startTime);
+      } else {
+        shaka.log.info(logPrefix, 'pop and abort prefetched init segment');
+      }
     }
   }
 
@@ -223,8 +279,8 @@ shaka.media.SegmentPrefetchOperation = class {
   /**
    * Fetch a segments
    *
-   * @param {!shaka.media.SegmentReference}
-   *   reference
+   * @param {!(shaka.media.SegmentReference|shaka.media.InitSegmentReference)}
+   *        reference
    * @param {!shaka.extern.Stream} stream
    * @public
    */

lib/media/streaming_engine.js

@@ -1180,6 +1180,12 @@ shaka.media.StreamingEngine = class {
     }
 
     if (mediaState.segmentPrefetch && mediaState.segmentIterator) {
+      const initSegmentReference = reference.initSegmentReference;
+      if (initSegmentReference && (!mediaState.lastSegmentReference ||
+          !shaka.media.InitSegmentReference.equal(
+              initSegmentReference, mediaState.lastInitSegmentReference))) {
+        mediaState.segmentPrefetch.prefetchInitSegment(initSegmentReference);
+      }
       mediaState.segmentPrefetch.prefetchSegments(reference);
     }
     // eslint-disable-next-line
@@ -1400,7 +1406,8 @@ shaka.media.StreamingEngine = class {
                   mediaState, presentationTime, stream, reference, dataToAppend);
 
             if (mediaState.segmentPrefetch && mediaState.segmentIterator) {
-              mediaState.segmentPrefetch.prefetchSegments(reference);
+              mediaState.segmentPrefetch.prefetchSegments(
+                  reference, /* skipFirst= */ true);
             }
     // eslint-disable-next-line
     Array.from(this.streamMediaStateMap_.values()).forEach(async (mediaState) => {
@@ -1445,7 +1452,8 @@ shaka.media.StreamingEngine = class {
             mediaState, presentationTime, stream, reference, mediaSegment);
 
         if (mediaState.segmentPrefetch && mediaState.segmentIterator) {
-          mediaState.segmentPrefetch.prefetchSegments(reference);
+          mediaState.segmentPrefetch.prefetchSegments(
+              reference, /* skipFirst= */ true);
         }
     // eslint-disable-next-line
     Array.from(this.streamMediaStateMap_.values()).forEach(async (mediaState) => {
@@ -2087,11 +2095,14 @@ shaka.media.StreamingEngine = class {
    * @suppress {strictMissingProperties}
    */
   async fetch_(mediaState, reference, streamDataCallback) {
+    if (reference instanceof shaka.media.InitSegmentReference) {
+      const segmentData = reference.getSegmentData();
+      if (segmentData) {
+        return segmentData;
+      }
+    }
     let op = null;
-    if (
-      mediaState.segmentPrefetch &&
-      reference instanceof shaka.media.SegmentReference
-    ) {
+    if (mediaState.segmentPrefetch) {
       op = mediaState.segmentPrefetch.getPrefetchedSegment(
           reference, streamDataCallback);
 
@@ -2101,8 +2112,7 @@ shaka.media.StreamingEngine = class {
     }
     if (!op) {
       op = this.dispatchFetch_(
-          reference, mediaState.stream, streamDataCallback,
-      );
+          reference, mediaState.stream, streamDataCallback);
     }
 
     mediaState.operation = op;
@@ -2124,7 +2134,6 @@ shaka.media.StreamingEngine = class {
    */
   dispatchFetch_(reference, stream, streamDataCallback) {
     const requestType = shaka.net.NetworkingEngine.RequestType.SEGMENT;
-
     const request = shaka.util.Networking.createSegmentRequest(
         reference.getUris(),
         reference.startByte,

test/test/util/fake_segment_prefetch.js

@@ -79,4 +79,7 @@ shaka.test.FakeSegmentPrefetch = class {
     }
     return null;
   }
+
+  /** @override */
+  prefetchInitSegment(reference) {}
 };