fix(youtube): restore SABR initialization cold start

extractor/src/main/java/org/schabi/newpipe/extractor/services/youtube/sabr/YoutubeSabrRequestBuilder.java

@@ -50,10 +50,13 @@ private static byte[] buildFirstMediaRequestLocked(@Nonnull final YoutubeSabrInf
         final List<SabrBufferedRange> bufferedRanges = streamState == null
                 ? java.util.Collections.emptyList()
                 : streamState.getBufferedRanges();
-        final boolean includeInitialPlaybackState = playerTimeMs > 0 || !bufferedRanges.isEmpty();
+        final boolean forcedInitialPlaybackState = streamState != null
+                && streamState.shouldWriteFirstRequestPlaybackState();
+        final boolean includeInitialPlaybackState = forcedInitialPlaybackState
+                || playerTimeMs > 0 || !bufferedRanges.isEmpty();
         final SabrProto.Writer request = new SabrProto.Writer();
         request.writeMessage(1, buildClientAbrState(audioFormat, videoFormat, playerTimeMs,
-                includeInitialPlaybackState,
+                includeInitialPlaybackState && !forcedInitialPlaybackState,
                 streamState == null
                         ? ENABLED_TRACK_TYPES_VIDEO_AND_AUDIO
                         : streamState.getEnabledTrackTypesBitfield(),
@@ -270,21 +273,21 @@ private static void writePreferredFormats(@Nonnull final SabrProto.Writer reques
                 return;
             }
             for (final YoutubeSabrFormat format : info.getFormats()) {
-                if (format.isAudio() && streamState.shouldSelectAudioFormat()) {
+                if (format.isAudio() && streamState.shouldPreferAudioFormat()) {
                     request.writeMessage(16, SabrProto.formatId(format));
                 }
             }
             for (final YoutubeSabrFormat format : info.getFormats()) {
-                if (format.isVideo() && streamState.shouldSelectVideoFormat()) {
+                if (format.isVideo() && streamState.shouldPreferVideoFormat()) {
                     request.writeMessage(17, SabrProto.formatId(format));
                 }
             }
             return;
         }
-        if (streamState == null || streamState.shouldSelectAudioFormat()) {
+        if (streamState == null || streamState.shouldPreferAudioFormat()) {
             request.writeMessage(16, SabrProto.formatId(audioFormat));
         }
-        if (streamState == null || streamState.shouldSelectVideoFormat()) {
+        if (streamState == null || streamState.shouldPreferVideoFormat()) {
             request.writeMessage(17, SabrProto.formatId(videoFormat));
         }
     }

extractor/src/main/java/org/schabi/newpipe/extractor/services/youtube/sabr/YoutubeSabrSession.java

@@ -182,6 +182,23 @@ public SabrMediaSegment fetchSegment(@Nonnull final SabrSegmentRequest request,
         if (cachedSegment != null) {
             return cachedSegment;
         }
+        final boolean initializationSegment = request.isInitializationSegment();
+        if (initializationSegment) {
+            prepareInitializationRequest(request.getFormat());
+        }
+        try {
+            return fetchUncachedSegment(request, localization);
+        } finally {
+            if (initializationSegment) {
+                clearInitializationRequest();
+            }
+        }
+    }
+
+    @Nonnull
+    private SabrMediaSegment fetchUncachedSegment(@Nonnull final SabrSegmentRequest request,
+                                                   @Nonnull final Localization localization)
+            throws IOException, ExtractionException {
         failIfKnownOutOfBounds(request);
 
         boolean targetPrepared = maybePrepareForDistantMediaSegment(request);
@@ -244,6 +261,28 @@ public SabrMediaSegment fetchSegment(@Nonnull final SabrSegmentRequest request,
                 : ":seq=" + request.getSequenceNumber()));
     }
 
+    private void prepareInitializationRequest(@Nonnull final YoutubeSabrFormat format) {
+        streamState.setWriteFirstRequestPlaybackState(true);
+        streamState.setWriteTopLevelPlayerTimeMs(false);
+        streamState.setWriteLastManualSelectedResolution(format.isVideo());
+        streamState.setBufferedRangesOverride(Collections.emptyList());
+        streamState.setRequestTrackMode(format.isVideo()
+                        ? YoutubeSabrStreamState.TRACK_MODE_VIDEO_ONLY
+                        : YoutubeSabrStreamState.TRACK_MODE_AUDIO_ONLY,
+                false, false);
+        streamState.setPreferredTrackTypes(true, true);
+        streamState.setPlayerTimeMs(streamState.getPlayerTimeMs());
+    }
+
+    private void clearInitializationRequest() {
+        streamState.setWriteFirstRequestPlaybackState(false);
+        streamState.setWriteTopLevelPlayerTimeMs(true);
+        streamState.setWriteLastManualSelectedResolution(false);
+        streamState.setBufferedRangesOverride(null);
+        streamState.clearPlayerTimeMsOverride();
+        streamState.setActiveTrackTypes(true, true);
+    }
+
     @Nonnull
     public YoutubeSabrProbeResult fetchNextResponse(@Nonnull final Localization localization)
             throws IOException, ExtractionException {

extractor/src/main/java/org/schabi/newpipe/extractor/services/youtube/sabr/YoutubeSabrStreamState.java

@@ -37,6 +37,9 @@ public final class YoutubeSabrStreamState {
     private volatile int enabledTrackTypesBitfield = YoutubeSabrRequestBuilder.ENABLED_TRACK_TYPES_VIDEO_AND_AUDIO;
     private volatile boolean selectAudioFormat = true;
     private volatile boolean selectVideoFormat = true;
+    private volatile boolean preferAudioFormat = true;
+    private volatile boolean preferVideoFormat = true;
+    private boolean writeFirstRequestPlaybackState;
     private boolean writeTopLevelPlayerTimeMs = true;
     private int clientViewportWidth = -1;
     private int clientViewportHeight = -1;
@@ -409,6 +412,14 @@ public synchronized void setRequestTrackMode(final int enabledTrackTypesBitfield
         this.enabledTrackTypesBitfield = enabledTrackTypesBitfield;
         this.selectAudioFormat = selectAudioFormat;
         this.selectVideoFormat = selectVideoFormat;
+        this.preferAudioFormat = selectAudioFormat;
+        this.preferVideoFormat = selectVideoFormat;
+    }
+
+    synchronized void setPreferredTrackTypes(final boolean videoActive,
+                                             final boolean audioActive) {
+        preferAudioFormat = audioActive;
+        preferVideoFormat = videoActive;
     }
 
     public void setActiveTrackTypes(final boolean videoActive, final boolean audioActive) {
@@ -495,6 +506,22 @@ boolean shouldSelectVideoFormat() {
         return selectVideoFormat;
     }
 
+    boolean shouldPreferAudioFormat() {
+        return preferAudioFormat;
+    }
+
+    boolean shouldPreferVideoFormat() {
+        return preferVideoFormat;
+    }
+
+    void setWriteFirstRequestPlaybackState(final boolean writeFirstRequestPlaybackState) {
+        this.writeFirstRequestPlaybackState = writeFirstRequestPlaybackState;
+    }
+
+    boolean shouldWriteFirstRequestPlaybackState() {
+        return writeFirstRequestPlaybackState;
+    }
+
     public void setWriteTopLevelPlayerTimeMs(final boolean writeTopLevelPlayerTimeMs) {
         this.writeTopLevelPlayerTimeMs = writeTopLevelPlayerTimeMs;
     }