Support multiple track outputs from BaseMediaChunk

Issue: #2362
Issue: #2176

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=148764237

library/src/main/java/com/google/android/exoplayer2/drm/OfflineLicenseHelper.java

@@ -177,8 +177,7 @@ public final class OfflineLicenseHelper<T extends ExoMediaCrypto> {
     Representation representation = adaptationSet.representations.get(0);
     DrmInitData drmInitData = representation.format.drmInitData;
     if (drmInitData == null) {
-      Format sampleFormat = DashUtil.loadSampleFormat(dataSource, representation,
-          adaptationSet.type);
+      Format sampleFormat = DashUtil.loadSampleFormat(dataSource, representation);
       if (sampleFormat != null) {
         drmInitData = sampleFormat.drmInitData;
       }

library/src/main/java/com/google/android/exoplayer2/source/chunk/BaseMediaChunk.java

@@ -21,14 +21,12 @@ import com.google.android.exoplayer2.upstream.DataSource;
 import com.google.android.exoplayer2.upstream.DataSpec;
 
 /**
- * A base implementation of {@link MediaChunk}, for chunks that contain a single track.
- * <p>
- * Loaded samples are output to a {@link DefaultTrackOutput}.
+ * A base implementation of {@link MediaChunk} that outputs to a {@link BaseMediaChunkOutput}.
  */
 public abstract class BaseMediaChunk extends MediaChunk {
 
-  private DefaultTrackOutput trackOutput;
-  private int firstSampleIndex;
+  private BaseMediaChunkOutput output;
+  private int[] firstSampleIndices;
 
   /**
    * @param dataSource The source from which the data should be loaded.
@@ -48,29 +46,29 @@ public abstract class BaseMediaChunk extends MediaChunk {
   }
 
   /**
-   * Initializes the chunk for loading, setting the {@link DefaultTrackOutput} that will receive
+   * Initializes the chunk for loading, setting the {@link BaseMediaChunkOutput} that will receive
    * samples as they are loaded.
    *
-   * @param trackOutput The output that will receive the loaded samples.
+   * @param output The output that will receive the loaded media samples.
    */
-  public void init(DefaultTrackOutput trackOutput) {
-    this.trackOutput = trackOutput;
-    this.firstSampleIndex = trackOutput.getWriteIndex();
+  public void init(BaseMediaChunkOutput output) {
+    this.output = output;
+    firstSampleIndices = output.getWriteIndices();
   }
 
   /**
-   * Returns the index of the first sample in the output that was passed to
-   * {@link #init(DefaultTrackOutput)} that will originate from this chunk.
+   * Returns the index of the first sample in the specified track of the output that will originate
+   * from this chunk.
    */
-  public final int getFirstSampleIndex() {
-    return firstSampleIndex;
+  public final int getFirstSampleIndex(int trackIndex) {
+    return firstSampleIndices[trackIndex];
   }
 
   /**
-   * Returns the track output most recently passed to {@link #init(DefaultTrackOutput)}.
+   * Returns the output most recently passed to {@link #init(BaseMediaChunkOutput)}.
    */
-  protected final DefaultTrackOutput getTrackOutput() {
-    return trackOutput;
+  protected final BaseMediaChunkOutput getOutput() {
+    return output;
   }
 
 }

library/src/main/java/com/google/android/exoplayer2/source/chunk/BaseMediaChunkOutput.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.android.exoplayer2.source.chunk;
+
+import android.util.Log;
+import com.google.android.exoplayer2.extractor.DefaultTrackOutput;
+import com.google.android.exoplayer2.extractor.DummyTrackOutput;
+import com.google.android.exoplayer2.extractor.TrackOutput;
+import com.google.android.exoplayer2.source.chunk.ChunkExtractorWrapper.TrackOutputProvider;
+
+/**
+ * An output for {@link BaseMediaChunk}s.
+ */
+/* package */ final class BaseMediaChunkOutput implements TrackOutputProvider {
+
+  private static final String TAG = "BaseMediaChunkOutput";
+
+  private final int[] trackTypes;
+  private final DefaultTrackOutput[] trackOutputs;
+
+  /**
+   * @param trackTypes The track types of the individual track outputs.
+   * @param trackOutputs The individual track outputs.
+   */
+  public BaseMediaChunkOutput(int[] trackTypes, DefaultTrackOutput... trackOutputs) {
+    this.trackTypes = trackTypes;
+    this.trackOutputs = trackOutputs;
+  }
+
+  @Override
+  public TrackOutput track(int id, int type) {
+    for (int i = 0; i < trackTypes.length; i++) {
+      if (type == trackTypes[i]) {
+        return trackOutputs[i];
+      }
+    }
+    Log.e(TAG, "Unmatched track of type: " + type);
+    return new DummyTrackOutput();
+  }
+
+  /**
+   * Returns the current absolute write indices of the individual track outputs.
+   */
+  public int[] getWriteIndices() {
+    int[] writeIndices = new int[trackOutputs.length];
+    for (int i = 0; i < trackOutputs.length; i++) {
+      if (trackOutputs[i] != null) {
+        writeIndices[i] = trackOutputs[i].getWriteIndex();
+      }
+    }
+    return writeIndices;
+  }
+
+  /**
+   * Sets an offset that will be added to the timestamps (and sub-sample timestamps) of samples
+   * subsequently written to the track outputs.
+   */
+  public void setSampleOffsetUs(long sampleOffsetUs) {
+    for (DefaultTrackOutput trackOutput : trackOutputs) {
+      if (trackOutput != null) {
+        trackOutput.setSampleOffsetUs(sampleOffsetUs);
+      }
+    }
+  }
+
+}

library/src/main/java/com/google/android/exoplayer2/source/chunk/ChunkExtractorWrapper.java

@@ -15,6 +15,7 @@
  */
 package com.google.android.exoplayer2.source.chunk;
 
+import android.util.SparseArray;
 import com.google.android.exoplayer2.C;
 import com.google.android.exoplayer2.Format;
 import com.google.android.exoplayer2.extractor.DummyTrackOutput;
@@ -32,33 +33,46 @@ import java.io.IOException;
  * <p>
  * The wrapper allows switching of the {@link TrackOutput} that receives parsed data.
  */
-public final class ChunkExtractorWrapper implements ExtractorOutput, TrackOutput {
+public final class ChunkExtractorWrapper implements ExtractorOutput {
+
+  /**
+   * Provides {@link TrackOutput} instances to be written to by the wrapper.
+   */
+  public interface TrackOutputProvider {
+
+    /**
+     * Called to get the {@link TrackOutput} for a specific track.
+     * <p>
+     * The same {@link TrackOutput} is returned if multiple calls are made with the same {@code id}.
+     *
+     * @param id A track identifier.
+     * @param type The type of the track. Typically one of the
+     *     {@link com.google.android.exoplayer2.C} {@code TRACK_TYPE_*} constants.
+     * @return The {@link TrackOutput} for the given track identifier.
+     */
+    TrackOutput track(int id, int type);
+
+  }
 
   public final Extractor extractor;
 
   private final Format manifestFormat;
-  private final int primaryTrackType;
+  private final SparseArray<BindingTrackOutput> bindingTrackOutputs;
 
   private boolean extractorInitialized;
-  private TrackOutput trackOutput;
+  private TrackOutputProvider trackOutputProvider;
   private SeekMap seekMap;
-  private Format sampleFormat;
-
-  // Accessed only on the loader thread.
-  private boolean seenTrack;
-  private int seenTrackId;
+  private Format[] sampleFormats;
 
   /**
    * @param extractor The extractor to wrap.
    * @param manifestFormat A manifest defined {@link Format} whose data should be merged into any
    *     sample {@link Format} output from the {@link Extractor}.
-   * @param primaryTrackType The type of the primary track. Typically one of the {@link C}
-   *     {@code TRACK_TYPE_*} constants.
    */
-  public ChunkExtractorWrapper(Extractor extractor, Format manifestFormat, int primaryTrackType) {
+  public ChunkExtractorWrapper(Extractor extractor, Format manifestFormat) {
     this.extractor = extractor;
     this.manifestFormat = manifestFormat;
-    this.primaryTrackType = primaryTrackType;
+    bindingTrackOutputs = new SparseArray<>();
   }
 
   /**
@@ -69,27 +83,27 @@ public final class ChunkExtractorWrapper implements ExtractorOutput, TrackOutput
   }
 
   /**
-   * Returns the sample {@link Format} most recently output by the extractor, or null.
+   * Returns the sample {@link Format}s most recently output by the extractor, or null.
    */
-  public Format getSampleFormat() {
-    return sampleFormat;
+  public Format[] getSampleFormats() {
+    return sampleFormats;
   }
 
   /**
    * Initializes the extractor to output to the provided {@link TrackOutput}, and configures it to
    * receive data from a new chunk.
    *
-   * @param trackOutput The {@link TrackOutput} that will receive sample data.
+   * @param trackOutputProvider The provider of {@link TrackOutput}s that will receive sample data.
    */
-  public void init(TrackOutput trackOutput) {
-    this.trackOutput = trackOutput;
+  public void init(TrackOutputProvider trackOutputProvider) {
+    this.trackOutputProvider = trackOutputProvider;
     if (!extractorInitialized) {
       extractor.init(this);
       extractorInitialized = true;
     } else {
       extractor.seek(0, 0);
-      if (sampleFormat != null && trackOutput != null) {
-        trackOutput.format(sampleFormat);
+      for (int i = 0; i < bindingTrackOutputs.size(); i++) {
+        bindingTrackOutputs.valueAt(i).bind(trackOutputProvider);
       }
     }
   }
@@ -98,18 +112,24 @@ public final class ChunkExtractorWrapper implements ExtractorOutput, TrackOutput
 
   @Override
   public TrackOutput track(int id, int type) {
-    if (primaryTrackType != C.TRACK_TYPE_UNKNOWN && primaryTrackType != type) {
-      return new DummyTrackOutput();
+    BindingTrackOutput bindingTrackOutput = bindingTrackOutputs.get(id);
+    if (bindingTrackOutput == null) {
+      // Assert that if we're seeing a new track we have not seen endTracks.
+      Assertions.checkState(sampleFormats == null);
+      bindingTrackOutput = new BindingTrackOutput(id, type, manifestFormat);
+      bindingTrackOutput.bind(trackOutputProvider);
+      bindingTrackOutputs.put(id, bindingTrackOutput);
     }
-    Assertions.checkState(!seenTrack || seenTrackId == id);
-    seenTrack = true;
-    seenTrackId = id;
-    return this;
+    return bindingTrackOutput;
   }
 
   @Override
   public void endTracks() {
-    Assertions.checkState(seenTrack);
+    Format[] sampleFormats = new Format[bindingTrackOutputs.size()];
+    for (int i = 0; i < bindingTrackOutputs.size(); i++) {
+      sampleFormats[i] = bindingTrackOutputs.valueAt(i).sampleFormat;
+    }
+    this.sampleFormats = sampleFormats;
   }
 
   @Override
@@ -117,31 +137,59 @@ public final class ChunkExtractorWrapper implements ExtractorOutput, TrackOutput
     this.seekMap = seekMap;
   }
 
-  // TrackOutput implementation.
+  // Internal logic.
 
-  @Override
-  public void format(Format format) {
-    sampleFormat = format.copyWithManifestFormatInfo(manifestFormat);
-    if (trackOutput != null) {
+  private static final class BindingTrackOutput implements TrackOutput {
+
+    private final int id;
+    private final int type;
+    private final Format manifestFormat;
+
+    public Format sampleFormat;
+    private TrackOutput trackOutput;
+
+    public BindingTrackOutput(int id, int type, Format manifestFormat) {
+      this.id = id;
+      this.type = type;
+      this.manifestFormat = manifestFormat;
+    }
+
+    public void bind(TrackOutputProvider trackOutputProvider) {
+      if (trackOutputProvider == null) {
+        trackOutput = new DummyTrackOutput();
+        return;
+      }
+      trackOutput = trackOutputProvider.track(id, type);
+      if (trackOutput != null) {
+        trackOutput.format(sampleFormat);
+      }
+    }
+
+    @Override
+    public void format(Format format) {
+      // TODO: This should only happen for the primary track. Additional metadata/text tracks need
+      // to be copied with different manifest derived formats.
+      sampleFormat = format.copyWithManifestFormatInfo(manifestFormat);
       trackOutput.format(sampleFormat);
     }
-  }
 
-  @Override
-  public int sampleData(ExtractorInput input, int length, boolean allowEndOfInput)
-      throws IOException, InterruptedException {
-    return trackOutput.sampleData(input, length, allowEndOfInput);
-  }
+    @Override
+    public int sampleData(ExtractorInput input, int length, boolean allowEndOfInput)
+        throws IOException, InterruptedException {
+      return trackOutput.sampleData(input, length, allowEndOfInput);
+    }
 
-  @Override
-  public void sampleData(ParsableByteArray data, int length) {
-    trackOutput.sampleData(data, length);
-  }
+    @Override
+    public void sampleData(ParsableByteArray data, int length) {
+      trackOutput.sampleData(data, length);
+    }
+
+    @Override
+    public void sampleMetadata(long timeUs, @C.BufferFlags int flags, int size, int offset,
+        byte[] encryptionKey) {
+      trackOutput.sampleMetadata(timeUs, flags, size, offset, encryptionKey);
+    }
 
-  @Override
-  public void sampleMetadata(long timeUs, @C.BufferFlags int flags, int size, int offset,
-      byte[] encryptionKey) {
-    trackOutput.sampleMetadata(timeUs, flags, size, offset, encryptionKey);
   }
 
 }

library/src/main/java/com/google/android/exoplayer2/source/chunk/ChunkSampleStream.java

@@ -44,6 +44,7 @@ public class ChunkSampleStream<T extends ChunkSource> implements SampleStream, S
   private final int minLoadableRetryCount;
   private final LinkedList<BaseMediaChunk> mediaChunks;
   private final List<BaseMediaChunk> readOnlyMediaChunks;
+  private final BaseMediaChunkOutput mediaChunkOutput;
   private final DefaultTrackOutput sampleQueue;
   private final ChunkHolder nextChunkHolder;
   private final Loader loader;
@@ -78,6 +79,7 @@ public class ChunkSampleStream<T extends ChunkSource> implements SampleStream, S
     mediaChunks = new LinkedList<>();
     readOnlyMediaChunks = Collections.unmodifiableList(mediaChunks);
     sampleQueue = new DefaultTrackOutput(allocator);
+    mediaChunkOutput = new BaseMediaChunkOutput(new int[] {trackType}, sampleQueue);
     lastSeekPositionUs = positionUs;
     pendingResetPositionUs = positionUs;
   }
@@ -127,7 +129,7 @@ public class ChunkSampleStream<T extends ChunkSource> implements SampleStream, S
     if (seekInsideBuffer) {
       // We succeeded. All we need to do is discard any chunks that we've moved past.
       while (mediaChunks.size() > 1
-          && mediaChunks.get(1).getFirstSampleIndex() <= sampleQueue.getReadIndex()) {
+          && mediaChunks.get(1).getFirstSampleIndex(0) <= sampleQueue.getReadIndex()) {
         mediaChunks.removeFirst();
       }
     } else {
@@ -176,7 +178,7 @@ public class ChunkSampleStream<T extends ChunkSource> implements SampleStream, S
     }
 
     while (mediaChunks.size() > 1
-        && mediaChunks.get(1).getFirstSampleIndex() <= sampleQueue.getReadIndex()) {
+        && mediaChunks.get(1).getFirstSampleIndex(0) <= sampleQueue.getReadIndex()) {
       mediaChunks.removeFirst();
     }
     BaseMediaChunk currentChunk = mediaChunks.getFirst();
@@ -232,7 +234,7 @@ public class ChunkSampleStream<T extends ChunkSource> implements SampleStream, S
       if (isMediaChunk) {
         BaseMediaChunk removed = mediaChunks.removeLast();
         Assertions.checkState(removed == loadable);
-        sampleQueue.discardUpstreamSamples(removed.getFirstSampleIndex());
+        sampleQueue.discardUpstreamSamples(removed.getFirstSampleIndex(0));
         if (mediaChunks.isEmpty()) {
           pendingResetPositionUs = lastSeekPositionUs;
         }
@@ -277,7 +279,7 @@ public class ChunkSampleStream<T extends ChunkSource> implements SampleStream, S
     if (isMediaChunk(loadable)) {
       pendingResetPositionUs = C.TIME_UNSET;
       BaseMediaChunk mediaChunk = (BaseMediaChunk) loadable;
-      mediaChunk.init(sampleQueue);
+      mediaChunk.init(mediaChunkOutput);
       mediaChunks.add(mediaChunk);
     }
     long elapsedRealtimeMs = loader.startLoading(loadable, this, minLoadableRetryCount);
@@ -337,7 +339,7 @@ public class ChunkSampleStream<T extends ChunkSource> implements SampleStream, S
       startTimeUs = removed.startTimeUs;
       loadingFinished = false;
     }
-    sampleQueue.discardUpstreamSamples(removed.getFirstSampleIndex());
+    sampleQueue.discardUpstreamSamples(removed.getFirstSampleIndex(0));
     eventDispatcher.upstreamDiscarded(trackType, startTimeUs, endTimeUs);
     return true;
   }

library/src/main/java/com/google/android/exoplayer2/source/chunk/ContainerMediaChunk.java

@@ -17,7 +17,6 @@ package com.google.android.exoplayer2.source.chunk;
 
 import com.google.android.exoplayer2.Format;
 import com.google.android.exoplayer2.extractor.DefaultExtractorInput;
-import com.google.android.exoplayer2.extractor.DefaultTrackOutput;
 import com.google.android.exoplayer2.extractor.Extractor;
 import com.google.android.exoplayer2.extractor.ExtractorInput;
 import com.google.android.exoplayer2.upstream.DataSource;
@@ -100,10 +99,10 @@ public class ContainerMediaChunk extends BaseMediaChunk {
       ExtractorInput input = new DefaultExtractorInput(dataSource,
           loadDataSpec.absoluteStreamPosition, dataSource.open(loadDataSpec));
       if (bytesLoaded == 0) {
-        // Set the target to ourselves.
-        DefaultTrackOutput trackOutput = getTrackOutput();
-        trackOutput.setSampleOffsetUs(sampleOffsetUs);
-        extractorWrapper.init(trackOutput);
+        // Configure the output and set it as the target for the extractor wrapper.
+        BaseMediaChunkOutput output = getOutput();
+        output.setSampleOffsetUs(sampleOffsetUs);
+        extractorWrapper.init(output);
       }
       // Load and decode the sample data.
       try {

library/src/main/java/com/google/android/exoplayer2/source/chunk/SingleSampleMediaChunk.java

@@ -18,8 +18,8 @@ package com.google.android.exoplayer2.source.chunk;
 import com.google.android.exoplayer2.C;
 import com.google.android.exoplayer2.Format;
 import com.google.android.exoplayer2.extractor.DefaultExtractorInput;
-import com.google.android.exoplayer2.extractor.DefaultTrackOutput;
 import com.google.android.exoplayer2.extractor.ExtractorInput;
+import com.google.android.exoplayer2.extractor.TrackOutput;
 import com.google.android.exoplayer2.upstream.DataSource;
 import com.google.android.exoplayer2.upstream.DataSpec;
 import com.google.android.exoplayer2.util.Util;
@@ -30,6 +30,7 @@ import java.io.IOException;
  */
 public final class SingleSampleMediaChunk extends BaseMediaChunk {
 
+  private final int trackType;
   private final Format sampleFormat;
 
   private volatile int bytesLoaded;
@@ -45,15 +46,20 @@ public final class SingleSampleMediaChunk extends BaseMediaChunk {
    * @param startTimeUs The start time of the media contained by the chunk, in microseconds.
    * @param endTimeUs The end time of the media contained by the chunk, in microseconds.
    * @param chunkIndex The index of the chunk.
+   * @param trackType The type of the chunk. Typically one of the {@link C} {@code TRACK_TYPE_*}
+   *     constants.
+   * @param sampleFormat The {@link Format} of the sample in the chunk.
    */
   public SingleSampleMediaChunk(DataSource dataSource, DataSpec dataSpec, Format trackFormat,
       int trackSelectionReason, Object trackSelectionData, long startTimeUs, long endTimeUs,
-      int chunkIndex, Format sampleFormat) {
+      int chunkIndex, int trackType, Format sampleFormat) {
     super(dataSource, dataSpec, trackFormat, trackSelectionReason, trackSelectionData, startTimeUs,
         endTimeUs, chunkIndex);
+    this.trackType = trackType;
     this.sampleFormat = sampleFormat;
   }
 
+
   @Override
   public boolean isLoadCompleted() {
     return loadCompleted;
@@ -87,8 +93,9 @@ public final class SingleSampleMediaChunk extends BaseMediaChunk {
         length += bytesLoaded;
       }
       ExtractorInput extractorInput = new DefaultExtractorInput(dataSource, bytesLoaded, length);
-      DefaultTrackOutput trackOutput = getTrackOutput();
-      trackOutput.setSampleOffsetUs(0);
+      BaseMediaChunkOutput output = getOutput();
+      output.setSampleOffsetUs(0);
+      TrackOutput trackOutput = output.track(0, trackType);
       trackOutput.format(sampleFormat);
       // Load the sample data.
       int result = 0;

library/src/main/java/com/google/android/exoplayer2/source/dash/DashUtil.java

@@ -68,17 +68,15 @@ public final class DashUtil {
    *
    * @param dataSource The source from which the data should be loaded.
    * @param representation The representation which initialization chunk belongs to.
-   * @param type The type of the primary track. Typically one of the {@link C} {@code TRACK_TYPE_*}
-   *     constants.
    * @return the sample {@link Format} of the given representation.
    * @throws IOException Thrown when there is an error while loading.
    * @throws InterruptedException Thrown if the thread was interrupted.
    */
-  public static Format loadSampleFormat(DataSource dataSource, Representation representation,
-      int type) throws IOException, InterruptedException {
+  public static Format loadSampleFormat(DataSource dataSource, Representation representation)
+      throws IOException, InterruptedException {
     ChunkExtractorWrapper extractorWrapper = loadInitializationData(dataSource, representation,
-        type, false);
-    return extractorWrapper == null ? null : extractorWrapper.getSampleFormat();
+        false);
+    return extractorWrapper == null ? null : extractorWrapper.getSampleFormats()[0];
   }
 
   /**
@@ -87,16 +85,14 @@ public final class DashUtil {
    *
    * @param dataSource The source from which the data should be loaded.
    * @param representation The representation which initialization chunk belongs to.
-   * @param type The type of the primary track. Typically one of the {@link C} {@code TRACK_TYPE_*}
-   *     constants.
    * @return {@link ChunkIndex} of the given representation.
    * @throws IOException Thrown when there is an error while loading.
    * @throws InterruptedException Thrown if the thread was interrupted.
    */
-  public static ChunkIndex loadChunkIndex(DataSource dataSource, Representation representation,
-      int type) throws IOException, InterruptedException {
+  public static ChunkIndex loadChunkIndex(DataSource dataSource, Representation representation)
+      throws IOException, InterruptedException {
     ChunkExtractorWrapper extractorWrapper = loadInitializationData(dataSource, representation,
-        type, true);
+        true);
     return extractorWrapper == null ? null : (ChunkIndex) extractorWrapper.getSeekMap();
   }
 
@@ -106,8 +102,6 @@ public final class DashUtil {
    *
    * @param dataSource The source from which the data should be loaded.
    * @param representation The representation which initialization chunk belongs to.
-   * @param type The type of the primary track. Typically one of the {@link C} {@code TRACK_TYPE_*}
-   *     constants.
    * @param loadIndex Whether to load index data too.
    * @return A {@link ChunkExtractorWrapper} for the {@code representation}, or null if no
    *     initialization or (if requested) index data exists.
@@ -115,13 +109,13 @@ public final class DashUtil {
    * @throws InterruptedException Thrown if the thread was interrupted.
    */
   private static ChunkExtractorWrapper loadInitializationData(DataSource dataSource,
-      Representation representation, int type, boolean loadIndex)
+      Representation representation, boolean loadIndex)
       throws IOException, InterruptedException {
     RangedUri initializationUri = representation.getInitializationUri();
     if (initializationUri == null) {
       return null;
     }
-    ChunkExtractorWrapper extractorWrapper = newWrappedExtractor(representation.format, type);
+    ChunkExtractorWrapper extractorWrapper = newWrappedExtractor(representation.format);
     RangedUri requestUri;
     if (loadIndex) {
       RangedUri indexUri = representation.getIndexUri();
@@ -153,12 +147,12 @@ public final class DashUtil {
     initializationChunk.load();
   }
 
-  private static ChunkExtractorWrapper newWrappedExtractor(Format format, int trackType) {
+  private static ChunkExtractorWrapper newWrappedExtractor(Format format) {
     String mimeType = format.containerMimeType;
     boolean isWebm = mimeType.startsWith(MimeTypes.VIDEO_WEBM)
         || mimeType.startsWith(MimeTypes.AUDIO_WEBM);
     Extractor extractor = isWebm ? new MatroskaExtractor() : new FragmentedMp4Extractor();
-    return new ChunkExtractorWrapper(extractor, format, trackType);
+    return new ChunkExtractorWrapper(extractor, format);
   }
 
   private DashUtil() {}

library/src/main/java/com/google/android/exoplayer2/source/dash/DefaultDashChunkSource.java

@@ -185,7 +185,7 @@ public class DefaultDashChunkSource implements DashChunkSource {
 
     RangedUri pendingInitializationUri = null;
     RangedUri pendingIndexUri = null;
-    if (representationHolder.extractorWrapper.getSampleFormat() == null) {
+    if (representationHolder.extractorWrapper.getSampleFormats() == null) {
       pendingInitializationUri = selectedRepresentation.getInitializationUri();
     }
     if (segmentIndex == null) {
@@ -343,7 +343,8 @@ public class DefaultDashChunkSource implements DashChunkSource {
       DataSpec dataSpec = new DataSpec(segmentUri.resolveUri(baseUrl),
           segmentUri.start, segmentUri.length, representation.getCacheKey());
       return new SingleSampleMediaChunk(dataSource, dataSpec, trackFormat, trackSelectionReason,
-          trackSelectionData, startTimeUs, endTimeUs, firstSegmentNum, trackFormat);
+          trackSelectionData, startTimeUs, endTimeUs, firstSegmentNum,
+          representationHolder.trackType, trackFormat);
     } else {
       int segmentCount = 1;
       for (int i = 1; i < maxSegmentCount; i++) {
@@ -370,6 +371,7 @@ public class DefaultDashChunkSource implements DashChunkSource {
 
   protected static final class RepresentationHolder {
 
+    public final int trackType;
     public final ChunkExtractorWrapper extractorWrapper;
 
     public Representation representation;
@@ -382,6 +384,7 @@ public class DefaultDashChunkSource implements DashChunkSource {
         boolean enableEventMessageTrack, boolean enableCea608Track, int trackType) {
       this.periodDurationUs = periodDurationUs;
       this.representation = representation;
+      this.trackType = trackType;
       String containerMimeType = representation.format.containerMimeType;
       if (mimeTypeIsRawText(containerMimeType)) {
         extractorWrapper = null;
@@ -403,7 +406,7 @@ public class DefaultDashChunkSource implements DashChunkSource {
         }
         // Prefer drmInitData obtained from the manifest over drmInitData obtained from the stream,
         // as per DASH IF Interoperability Recommendations V3.0, 7.5.3.
-        extractorWrapper = new ChunkExtractorWrapper(extractor, representation.format, trackType);
+        extractorWrapper = new ChunkExtractorWrapper(extractor, representation.format);
       }
       segmentIndex = representation.getIndex();
     }

library/src/main/java/com/google/android/exoplayer2/source/smoothstreaming/DefaultSsChunkSource.java

@@ -102,7 +102,7 @@ public class DefaultSsChunkSource implements SsChunkSource {
       FragmentedMp4Extractor extractor = new FragmentedMp4Extractor(
           FragmentedMp4Extractor.FLAG_WORKAROUND_EVERY_VIDEO_FRAME_IS_SYNC_FRAME
           | FragmentedMp4Extractor.FLAG_WORKAROUND_IGNORE_TFDT_BOX, null, track);
-      extractorWrappers[i] = new ChunkExtractorWrapper(extractor, format, streamElement.type);
+      extractorWrappers[i] = new ChunkExtractorWrapper(extractor, format);
     }
   }