Add reference counting to DrmSession

This CL should not introduce any functional changes.

PiperOrigin-RevId: 250277165

extensions/vp9/src/main/java/com/google/android/exoplayer2/ext/vp9/LibvpxVideoRenderer.java

@@ -473,21 +473,13 @@ public class LibvpxVideoRenderer extends BaseRenderer {
   }
 
   private void setSourceDrmSession(@Nullable DrmSession<ExoMediaCrypto> session) {
-    DrmSession<ExoMediaCrypto> previous = sourceDrmSession;
+    DrmSession.replaceSessionReferences(sourceDrmSession, session);
     sourceDrmSession = session;
-    releaseDrmSessionIfUnused(previous);
   }
 
   private void setDecoderDrmSession(@Nullable DrmSession<ExoMediaCrypto> session) {
-    DrmSession<ExoMediaCrypto> previous = decoderDrmSession;
+    DrmSession.replaceSessionReferences(decoderDrmSession, session);
     decoderDrmSession = session;
-    releaseDrmSessionIfUnused(previous);
-  }
-
-  private void releaseDrmSessionIfUnused(@Nullable DrmSession<ExoMediaCrypto> session) {
-    if (session != null && session != decoderDrmSession && session != sourceDrmSession) {
-      drmSessionManager.releaseSession(session);
-    }
   }
 
   /**
@@ -512,12 +504,10 @@ public class LibvpxVideoRenderer extends BaseRenderer {
         }
         DrmSession<ExoMediaCrypto> session =
             drmSessionManager.acquireSession(Looper.myLooper(), newFormat.drmInitData);
-        if (session == decoderDrmSession || session == sourceDrmSession) {
-          // We already had this session. The manager must be reference counting, so release it once
-          // to get the count attributed to this renderer back down to 1.
-          drmSessionManager.releaseSession(session);
+        if (sourceDrmSession != null) {
+          sourceDrmSession.releaseReference();
         }
-        setSourceDrmSession(session);
+        sourceDrmSession = session;
       } else {
         setSourceDrmSession(null);
       }

library/core/src/main/java/com/google/android/exoplayer2/FormatHolder.java

@@ -16,7 +16,7 @@
 package com.google.android.exoplayer2;
 
 import androidx.annotation.Nullable;
-import com.google.android.exoplayer2.drm.DecryptionResource;
+import com.google.android.exoplayer2.drm.DrmSession;
 
 /**
  * Holds a {@link Format}.
@@ -25,14 +25,14 @@ public final class FormatHolder {
 
   /**
    * Whether the object expected to populate {@link #format} is also expected to populate {@link
-   * #decryptionResource}.
+   * #drmSession}.
    */
   // TODO: Remove once all Renderers and MediaSources have migrated to the new DRM model [Internal
   // ref: b/129764794].
   public boolean decryptionResourceIsProvided;
 
   /** An accompanying context for decrypting samples in the format. */
-  @Nullable public DecryptionResource<?> decryptionResource;
+  @Nullable public DrmSession<?> drmSession;
 
   /** The held {@link Format}. */
   @Nullable public Format format;

library/core/src/main/java/com/google/android/exoplayer2/audio/SimpleDecoderAudioRenderer.java

@@ -646,21 +646,13 @@ public abstract class SimpleDecoderAudioRenderer extends BaseRenderer implements
   }
 
   private void setSourceDrmSession(@Nullable DrmSession<ExoMediaCrypto> session) {
-    DrmSession<ExoMediaCrypto> previous = sourceDrmSession;
+    DrmSession.replaceSessionReferences(sourceDrmSession, session);
     sourceDrmSession = session;
-    releaseDrmSessionIfUnused(previous);
   }
 
   private void setDecoderDrmSession(@Nullable DrmSession<ExoMediaCrypto> session) {
-    DrmSession<ExoMediaCrypto> previous = decoderDrmSession;
+    DrmSession.replaceSessionReferences(decoderDrmSession, session);
     decoderDrmSession = session;
-    releaseDrmSessionIfUnused(previous);
-  }
-
-  private void releaseDrmSessionIfUnused(@Nullable DrmSession<ExoMediaCrypto> session) {
-    if (session != null && session != decoderDrmSession && session != sourceDrmSession) {
-      drmSessionManager.releaseSession(session);
-    }
   }
 
   private void onInputFormatChanged(Format newFormat) throws ExoPlaybackException {
@@ -677,12 +669,10 @@ public abstract class SimpleDecoderAudioRenderer extends BaseRenderer implements
         }
         DrmSession<ExoMediaCrypto> session =
             drmSessionManager.acquireSession(Looper.myLooper(), newFormat.drmInitData);
-        if (session == decoderDrmSession || session == sourceDrmSession) {
-          // We already had this session. The manager must be reference counting, so release it once
-          // to get the count attributed to this renderer back down to 1.
-          drmSessionManager.releaseSession(session);
+        if (sourceDrmSession != null) {
+          sourceDrmSession.releaseReference();
         }
-        setSourceDrmSession(session);
+        sourceDrmSession = session;
       } else {
         setSourceDrmSession(null);
       }

library/core/src/main/java/com/google/android/exoplayer2/drm/DefaultDrmSession.java

@@ -38,6 +38,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import org.checkerframework.checker.nullness.compatqual.NullableType;
 import org.checkerframework.checker.nullness.qual.EnsuresNonNullIf;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 import org.checkerframework.checker.nullness.qual.RequiresNonNull;
@@ -103,12 +104,12 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
   /* package */ final PostResponseHandler postResponseHandler;
 
   private @DrmSession.State int state;
-  private int openCount;
+  private int referenceCount;
   @Nullable private HandlerThread requestHandlerThread;
   @Nullable private PostRequestHandler postRequestHandler;
   @Nullable private T mediaCrypto;
   @Nullable private DrmSessionException lastException;
-  private byte @MonotonicNonNull [] sessionId;
+  private byte @NullableType [] sessionId;
   private byte @MonotonicNonNull [] offlineLicenseKeySetId;
 
   @Nullable private KeyRequest currentKeyRequest;
@@ -169,42 +170,6 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
     postResponseHandler = new PostResponseHandler(playbackLooper);
   }
 
-  // Life cycle.
-
-  public void acquire() {
-    if (++openCount == 1) {
-      requestHandlerThread = new HandlerThread("DrmRequestHandler");
-      requestHandlerThread.start();
-      postRequestHandler = new PostRequestHandler(requestHandlerThread.getLooper());
-      if (openInternal(true)) {
-        doLicense(true);
-      }
-    }
-  }
-
-  @SuppressWarnings("assignment.type.incompatible")
-  public void release() {
-    if (--openCount == 0) {
-      // Assigning null to various non-null variables for clean-up.
-      state = STATE_RELEASED;
-      postResponseHandler.removeCallbacksAndMessages(null);
-      Util.castNonNull(postRequestHandler).removeCallbacksAndMessages(null);
-      postRequestHandler = null;
-      Util.castNonNull(requestHandlerThread).quit();
-      requestHandlerThread = null;
-      mediaCrypto = null;
-      lastException = null;
-      currentKeyRequest = null;
-      currentProvisionRequest = null;
-      if (sessionId != null) {
-        mediaDrm.closeSession(sessionId);
-        sessionId = null;
-        eventDispatcher.dispatch(DefaultDrmSessionEventListener::onDrmSessionReleased);
-      }
-      releaseCallback.onSessionReleased(this);
-    }
-  }
-
   public boolean hasSessionId(byte[] sessionId) {
     return Arrays.equals(this.sessionId, sessionId);
   }
@@ -270,6 +235,42 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
     return offlineLicenseKeySetId;
   }
 
+  @Override
+  public void acquireReference() {
+    if (++referenceCount == 1) {
+      Assertions.checkState(state == STATE_OPENING);
+      requestHandlerThread = new HandlerThread("DrmRequestHandler");
+      requestHandlerThread.start();
+      postRequestHandler = new PostRequestHandler(requestHandlerThread.getLooper());
+      if (openInternal(true)) {
+        doLicense(true);
+      }
+    }
+  }
+
+  @Override
+  public void releaseReference() {
+    if (--referenceCount == 0) {
+      // Assigning null to various non-null variables for clean-up.
+      state = STATE_RELEASED;
+      Util.castNonNull(postResponseHandler).removeCallbacksAndMessages(null);
+      Util.castNonNull(postRequestHandler).removeCallbacksAndMessages(null);
+      postRequestHandler = null;
+      Util.castNonNull(requestHandlerThread).quit();
+      requestHandlerThread = null;
+      mediaCrypto = null;
+      lastException = null;
+      currentKeyRequest = null;
+      currentProvisionRequest = null;
+      if (sessionId != null) {
+        mediaDrm.closeSession(sessionId);
+        sessionId = null;
+        eventDispatcher.dispatch(DefaultDrmSessionEventListener::onDrmSessionReleased);
+      }
+      releaseCallback.onSessionReleased(this);
+    }
+  }
+
   // Internal methods.
 
   /**
@@ -288,9 +289,10 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
 
     try {
       sessionId = mediaDrm.openSession();
-      eventDispatcher.dispatch(DefaultDrmSessionEventListener::onDrmSessionAcquired);
       mediaCrypto = mediaDrm.createMediaCrypto(sessionId);
+      eventDispatcher.dispatch(DefaultDrmSessionEventListener::onDrmSessionAcquired);
       state = STATE_OPENED;
+      Assertions.checkNotNull(sessionId);
       return true;
     } catch (NotProvisionedException e) {
       if (allowProvisioning) {
@@ -329,6 +331,7 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
 
   @RequiresNonNull("sessionId")
   private void doLicense(boolean allowRetry) {
+    byte[] sessionId = Util.castNonNull(this.sessionId);
     switch (mode) {
       case DefaultDrmSessionManager.MODE_PLAYBACK:
       case DefaultDrmSessionManager.MODE_QUERY:
@@ -364,6 +367,7 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
         break;
       case DefaultDrmSessionManager.MODE_RELEASE:
         Assertions.checkNotNull(offlineLicenseKeySetId);
+        Assertions.checkNotNull(this.sessionId);
         // It's not necessary to restore the key (and open a session to do that) before releasing it
         // but this serves as a good sanity/fast-failure check.
         if (restoreKeys()) {

library/core/src/main/java/com/google/android/exoplayer2/drm/DefaultDrmSessionManager.java

@@ -432,19 +432,10 @@ public class DefaultDrmSessionManager<T extends ExoMediaCrypto>
               initialDrmRequestRetryCount);
       sessions.add(session);
     }
-    session.acquire();
+    session.acquireReference();
     return session;
   }
 
-  @Override
-  public void releaseSession(DrmSession<T> session) {
-    if (session instanceof ErrorStateDrmSession) {
-      // Do nothing.
-      return;
-    }
-    ((DefaultDrmSession<T>) session).release();
-  }
-
   // ProvisioningManager implementation.
 
   @Override

library/core/src/main/java/com/google/android/exoplayer2/drm/DrmSession.java

@@ -28,6 +28,20 @@ import java.util.Map;
  */
 public interface DrmSession<T extends ExoMediaCrypto> {
 
+  /**
+   * Invokes {@code newSession's} {@link #acquireReference()} and {@code previousSession's} {@link
+   * #releaseReference()} in that order. Does nothing for passed null values.
+   */
+  static <T extends ExoMediaCrypto> void replaceSessionReferences(
+      @Nullable DrmSession<T> previousSession, @Nullable DrmSession<T> newSession) {
+    if (newSession != null) {
+      newSession.acquireReference();
+    }
+    if (previousSession != null) {
+      previousSession.releaseReference();
+    }
+  }
+
   /**
    * Wraps the throwable which is the cause of the error state.
    */
@@ -110,4 +124,18 @@ public interface DrmSession<T extends ExoMediaCrypto> {
    */
   @Nullable
   byte[] getOfflineLicenseKeySetId();
+
+  /**
+   * Increments the reference count for this session. A non-zero reference count session will keep
+   * any acquired resources.
+   */
+  void acquireReference();
+
+  /**
+   * Decreases by one the reference count for this session. A session that reaches a zero reference
+   * count will release any resources it holds.
+   *
+   * <p>The session must not be used after its reference count has been reduced to 0.
+   */
+  void releaseReference();
 }

library/core/src/main/java/com/google/android/exoplayer2/drm/DrmSessionManager.java

@@ -34,8 +34,10 @@ public interface DrmSessionManager<T extends ExoMediaCrypto> {
   boolean canAcquireSession(DrmInitData drmInitData);
 
   /**
-   * Acquires a {@link DrmSession} for the specified {@link DrmInitData}. The {@link DrmSession}
-   * must be returned to {@link #releaseSession(DrmSession)} when it is no longer required.
+   * Returns a {@link DrmSession} with an acquired reference for the specified {@link DrmInitData}.
+   *
+   * <p>The caller must call {@link DrmSession#releaseReference} to decrement the session's
+   * reference count when the session is no longer required.
    *
    * @param playbackLooper The looper associated with the media playback thread.
    * @param drmInitData DRM initialization data. All contained {@link SchemeData}s must contain
@@ -43,10 +45,4 @@ public interface DrmSessionManager<T extends ExoMediaCrypto> {
    * @return The DRM session.
    */
   DrmSession<T> acquireSession(Looper playbackLooper, DrmInitData drmInitData);
-
-  /**
-   * Releases a {@link DrmSession}.
-   */
-  void releaseSession(DrmSession<T> drmSession);
-
 }

library/core/src/main/java/com/google/android/exoplayer2/drm/ErrorStateDrmSession.java

@@ -57,4 +57,13 @@ public final class ErrorStateDrmSession<T extends ExoMediaCrypto> implements Drm
     return null;
   }
 
+  @Override
+  public void acquireReference() {
+    // Do nothing.
+  }
+
+  @Override
+  public void releaseReference() {
+    // Do nothing.
+  }
 }

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

@@ -235,7 +235,7 @@ public final class OfflineLicenseHelper<T extends ExoMediaCrypto> {
     DrmSessionException error = drmSession.getError();
     Pair<Long, Long> licenseDurationRemainingSec =
         WidevineUtil.getLicenseDurationRemainingSec(drmSession);
-    drmSessionManager.releaseSession(drmSession);
+    drmSession.releaseReference();
     if (error != null) {
       if (error.getCause() instanceof KeysExpiredException) {
         return Pair.create(0L, 0L);
@@ -259,7 +259,7 @@ public final class OfflineLicenseHelper<T extends ExoMediaCrypto> {
         drmInitData);
     DrmSessionException error = drmSession.getError();
     byte[] keySetId = drmSession.getOfflineLicenseKeySetId();
-    drmSessionManager.releaseSession(drmSession);
+    drmSession.releaseReference();
     if (error != null) {
       throw error;
     }

library/core/src/main/java/com/google/android/exoplayer2/mediacodec/MediaCodecRenderer.java

@@ -941,21 +941,13 @@ public abstract class MediaCodecRenderer extends BaseRenderer {
   }
 
   private void setSourceDrmSession(@Nullable DrmSession<FrameworkMediaCrypto> session) {
-    DrmSession<FrameworkMediaCrypto> previous = sourceDrmSession;
+    DrmSession.replaceSessionReferences(sourceDrmSession, session);
     sourceDrmSession = session;
-    releaseDrmSessionIfUnused(previous);
   }
 
   private void setCodecDrmSession(@Nullable DrmSession<FrameworkMediaCrypto> session) {
-    DrmSession<FrameworkMediaCrypto> previous = codecDrmSession;
+    DrmSession.replaceSessionReferences(codecDrmSession, session);
     codecDrmSession = session;
-    releaseDrmSessionIfUnused(previous);
-  }
-
-  private void releaseDrmSessionIfUnused(@Nullable DrmSession<FrameworkMediaCrypto> session) {
-    if (session != null && session != sourceDrmSession && session != codecDrmSession) {
-      drmSessionManager.releaseSession(session);
-    }
   }
 
   /**
@@ -1159,12 +1151,10 @@ public abstract class MediaCodecRenderer extends BaseRenderer {
         }
         DrmSession<FrameworkMediaCrypto> session =
             drmSessionManager.acquireSession(Looper.myLooper(), newFormat.drmInitData);
-        if (session == sourceDrmSession || session == codecDrmSession) {
-          // We already had this session. The manager must be reference counting, so release it once
-          // to get the count attributed to this renderer back down to 1.
-          drmSessionManager.releaseSession(session);
+        if (sourceDrmSession != null) {
+          sourceDrmSession.releaseReference();
         }
-        setSourceDrmSession(session);
+        sourceDrmSession = session;
       } else {
         setSourceDrmSession(null);
       }