From 87c4d60c37470732e16e2d917d57a9c4825817db Mon Sep 17 00:00:00 2001
From: Googler <noreply@google.com>
Date: Wed, 28 Feb 2024 10:24:32 -0800
Subject: [PATCH] Allow CronetDataSource's read buffer size to be configured.

PiperOrigin-RevId: 611146397
---
 .../datasource/cronet/CronetDataSource.java   | 33 +++++++++++++++----
 .../cronet/CronetDataSourceFactory.java       |  3 +-
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/libraries/datasource_cronet/src/main/java/androidx/media3/datasource/cronet/CronetDataSource.java b/libraries/datasource_cronet/src/main/java/androidx/media3/datasource/cronet/CronetDataSource.java
index 11376d9cbe4..d030564688d 100644
--- a/libraries/datasource_cronet/src/main/java/androidx/media3/datasource/cronet/CronetDataSource.java
+++ b/libraries/datasource_cronet/src/main/java/androidx/media3/datasource/cronet/CronetDataSource.java
@@ -93,6 +93,7 @@ public static final class Factory implements HttpDataSource.Factory {
     private int requestPriority;
     private int connectTimeoutMs;
     private int readTimeoutMs;
+    private int readBufferSize;
     private boolean resetTimeoutOnRedirects;
     private boolean handleSetCookieRequests;
     private boolean keepPostFor302Redirects;
@@ -117,6 +118,7 @@ public Factory(CronetEngine cronetEngine, Executor executor) {
       requestPriority = REQUEST_PRIORITY_MEDIUM;
       connectTimeoutMs = DEFAULT_CONNECT_TIMEOUT_MILLIS;
       readTimeoutMs = DEFAULT_READ_TIMEOUT_MILLIS;
+      readBufferSize = DEFAULT_READ_BUFFER_SIZE_BYTES;
     }
 
     /**
@@ -142,6 +144,7 @@ public Factory(CronetEngineWrapper cronetEngineWrapper, Executor executor) {
       internalFallbackFactory = new DefaultHttpDataSource.Factory();
       connectTimeoutMs = DEFAULT_CONNECT_TIMEOUT_MILLIS;
       readTimeoutMs = DEFAULT_READ_TIMEOUT_MILLIS;
+      readBufferSize = DEFAULT_READ_BUFFER_SIZE_BYTES;
     }
 
     @CanIgnoreReturnValue
@@ -334,6 +337,19 @@ public Factory setFallbackFactory(@Nullable HttpDataSource.Factory fallbackFacto
       return this;
     }
 
+    /**
+     * Sets the read buffer size, in bytes.
+     *
+     * @param readBufferSize The read buffer size, in bytes.
+     * @return This factory.
+     */
+    @CanIgnoreReturnValue
+    @UnstableApi
+    public Factory setReadBufferSize(int readBufferSize) {
+      this.readBufferSize = readBufferSize;
+      return this;
+    }
+
     @UnstableApi
     @Override
     public HttpDataSource createDataSource() {
@@ -354,7 +370,8 @@ public HttpDataSource createDataSource() {
               userAgent,
               defaultRequestProperties,
               contentTypePredicate,
-              keepPostFor302Redirects);
+              keepPostFor302Redirects,
+              readBufferSize);
       if (transferListener != null) {
         dataSource.addTransferListener(transferListener);
       }
@@ -421,10 +438,11 @@ public OpenException(
   /** The default read timeout, in milliseconds. */
   @UnstableApi public static final int DEFAULT_READ_TIMEOUT_MILLIS = 8 * 1000;
 
-  /* package */ final UrlRequest.Callback urlRequestCallback;
-
   // The size of read buffer passed to cronet UrlRequest.read().
-  private static final int READ_BUFFER_SIZE_BYTES = 32 * 1024;
+  // TODO: Make this private once CronetDataSourceFactory is deleted.
+  static final int DEFAULT_READ_BUFFER_SIZE_BYTES = 32 * 1024;
+
+  /* package */ final UrlRequest.Callback urlRequestCallback;
 
   private final CronetEngine cronetEngine;
   private final Executor executor;
@@ -438,6 +456,7 @@ public OpenException(
   private final RequestProperties requestProperties;
   private final ConditionVariable operation;
   private final Clock clock;
+  private final int readBufferSize;
 
   @Nullable private Predicate<String> contentTypePredicate;
   private final boolean keepPostFor302Redirects;
@@ -476,7 +495,8 @@ protected CronetDataSource(
       @Nullable String userAgent,
       @Nullable RequestProperties defaultRequestProperties,
       @Nullable Predicate<String> contentTypePredicate,
-      boolean keepPostFor302Redirects) {
+      boolean keepPostFor302Redirects,
+      int readBufferSize) {
     super(/* isNetwork= */ true);
     this.cronetEngine = Assertions.checkNotNull(cronetEngine);
     this.executor = Assertions.checkNotNull(executor);
@@ -490,6 +510,7 @@ protected CronetDataSource(
     this.contentTypePredicate = contentTypePredicate;
     this.keepPostFor302Redirects = keepPostFor302Redirects;
     clock = Clock.DEFAULT;
+    this.readBufferSize = readBufferSize;
     urlRequestCallback = new UrlRequestCallback();
     requestProperties = new RequestProperties();
     operation = new ConditionVariable();
@@ -1025,7 +1046,7 @@ private void readInternal(ByteBuffer buffer, DataSpec dataSpec) throws HttpDataS
 
   private ByteBuffer getOrCreateReadBuffer() {
     if (readBuffer == null) {
-      readBuffer = ByteBuffer.allocateDirect(READ_BUFFER_SIZE_BYTES);
+      readBuffer = ByteBuffer.allocateDirect(readBufferSize);
       readBuffer.limit(0);
     }
     return readBuffer;
diff --git a/libraries/datasource_cronet/src/main/java/androidx/media3/datasource/cronet/CronetDataSourceFactory.java b/libraries/datasource_cronet/src/main/java/androidx/media3/datasource/cronet/CronetDataSourceFactory.java
index b38c35ddfbc..827225422d6 100644
--- a/libraries/datasource_cronet/src/main/java/androidx/media3/datasource/cronet/CronetDataSourceFactory.java
+++ b/libraries/datasource_cronet/src/main/java/androidx/media3/datasource/cronet/CronetDataSourceFactory.java
@@ -370,7 +370,8 @@ protected HttpDataSource createDataSourceInternal(
             /* userAgent= */ null,
             defaultRequestProperties,
             /* contentTypePredicate= */ null,
-            /* keepPostFor302Redirects= */ false);
+            /* keepPostFor302Redirects= */ false,
+            CronetDataSource.DEFAULT_READ_BUFFER_SIZE_BYTES);
     if (transferListener != null) {
       dataSource.addTransferListener(transferListener);
     }