cronet: Add internal API to specify Network

ejona86 · ejona86 · commit c2eccca3bc10 · 2024-08-13T11:22:01.000-07:00
cl/661194496
diff --git a/cronet/src/main/java/io/grpc/cronet/CronetChannelBuilder.java b/cronet/src/main/java/io/grpc/cronet/CronetChannelBuilder.java
@@ -20,9 +20,12 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 import static io.grpc.internal.GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE;
 
+import android.net.Network;
+import android.os.Build;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.DoNotCall;
 import io.grpc.ChannelCredentials;
 import io.grpc.ChannelLogger;
@@ -105,6 +108,7 @@ public static CronetChannelBuilder forAddress(String name, int port) {
   private int trafficStatsTag;
   private boolean trafficStatsUidSet;
   private int trafficStatsUid;
+  private Network network;
 
   private CronetChannelBuilder(String host, int port, CronetEngine cronetEngine) {
     final class CronetChannelTransportFactoryBuilder implements ClientTransportFactoryBuilder {
@@ -190,6 +194,13 @@ CronetChannelBuilder setTrafficStatsUid(int uid) {
     return this;
   }
 
+  /** Sets the network ID to use for this channel traffic. */
+  @CanIgnoreReturnValue
+  CronetChannelBuilder bindToNetwork(@Nullable Network network) {
+    this.network = network;
+    return this;
+  }
+
   /**
    * Provides a custom scheduled executor service.
    *
@@ -210,7 +221,12 @@ public CronetChannelBuilder scheduledExecutorService(
   ClientTransportFactory buildTransportFactory() {
     return new CronetTransportFactory(
         new TaggingStreamFactory(
-            cronetEngine, trafficStatsTagSet, trafficStatsTag, trafficStatsUidSet, trafficStatsUid),
+            cronetEngine,
+            trafficStatsTagSet,
+            trafficStatsTag,
+            trafficStatsUidSet,
+            trafficStatsUid,
+            network),
         MoreExecutors.directExecutor(),
         scheduledExecutorService,
         maxMessageSize,
@@ -294,18 +310,21 @@ private static class TaggingStreamFactory extends StreamBuilderFactory {
     private final int trafficStatsTag;
     private final boolean trafficStatsUidSet;
     private final int trafficStatsUid;
+    private final Network network;
 
     TaggingStreamFactory(
         CronetEngine cronetEngine,
         boolean trafficStatsTagSet,
         int trafficStatsTag,
         boolean trafficStatsUidSet,
-        int trafficStatsUid) {
+        int trafficStatsUid,
+        Network network) {
       this.cronetEngine = cronetEngine;
       this.trafficStatsTagSet = trafficStatsTagSet;
       this.trafficStatsTag = trafficStatsTag;
       this.trafficStatsUidSet = trafficStatsUidSet;
       this.trafficStatsUid = trafficStatsUid;
+      this.network = network;
     }
 
     @Override
@@ -320,6 +339,11 @@ public BidirectionalStream.Builder newBidirectionalStreamBuilder(
       if (trafficStatsUidSet) {
         builder.setTrafficStatsUid(trafficStatsUid);
       }
+      if (network != null) {
+        if (Build.VERSION.SDK_INT >= 23) {
+          builder.bindToNetwork(network.getNetworkHandle());
+        }
+      }
       return builder;
     }
   }
diff --git a/cronet/src/main/java/io/grpc/cronet/InternalCronetChannelBuilder.java b/cronet/src/main/java/io/grpc/cronet/InternalCronetChannelBuilder.java
@@ -16,7 +16,9 @@
 
 package io.grpc.cronet;
 
+import android.net.Network;
 import io.grpc.Internal;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * Internal {@link CronetChannelBuilder} accessor. This is intended for usage internal to the gRPC
@@ -58,4 +60,9 @@ public static void setTrafficStatsTag(CronetChannelBuilder builder, int tag) {
   public static void setTrafficStatsUid(CronetChannelBuilder builder, int uid) {
     builder.setTrafficStatsUid(uid);
   }
+
+  /** Sets the network {@link android.net.Network} to use when relying traffic by this channel. */
+  public static void bindToNetwork(CronetChannelBuilder builder, @Nullable Network network) {
+    builder.bindToNetwork(network);
+  }
 }
