2 files changed, 8 insertions, 7 deletions
diff --git a/services/net/java/android/net/IpMemoryStore.java b/services/net/java/android/net/IpMemoryStore.java
index dcefb537d0a0..8df2e0d08e0e 100644
--- a/services/net/java/android/net/IpMemoryStore.java
+++ b/services/net/java/android/net/IpMemoryStore.java
@@ -18,6 +18,7 @@ package android.net;
 
 import android.annotation.NonNull;
 import android.content.Context;
+import android.net.networkstack.ModuleNetworkStackClient;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -41,7 +42,7 @@ public class IpMemoryStore extends IpMemoryStoreClient {
         super(context);
         mService = new CompletableFuture<>();
         mTailNode = new AtomicReference<CompletableFuture<IIpMemoryStore>>(mService);
-        getNetworkStackClient().fetchIpMemoryStore(
+        getModuleNetworkStackClient(context).fetchIpMemoryStore(
                 new IIpMemoryStoreCallbacks.Stub() {
                     @Override
                     public void onIpMemoryStoreFetched(@NonNull final IIpMemoryStore memoryStore) {
@@ -85,8 +86,8 @@ public class IpMemoryStore extends IpMemoryStoreClient {
     }
 
     @VisibleForTesting
-    protected NetworkStackClient getNetworkStackClient() {
-        return NetworkStackClient.getInstance();
+    protected ModuleNetworkStackClient getModuleNetworkStackClient(Context context) {
+        return ModuleNetworkStackClient.getInstance(context);
     }
 
     /** Gets an instance of the memory store */
diff --git a/services/net/java/android/net/ip/IpClientUtil.java b/services/net/java/android/net/ip/IpClientUtil.java
index a3618b47171a..b329aeec4853 100644
--- a/services/net/java/android/net/ip/IpClientUtil.java
+++ b/services/net/java/android/net/ip/IpClientUtil.java
@@ -22,7 +22,7 @@ import android.content.Context;
 import android.net.DhcpResultsParcelable;
 import android.net.Layer2PacketParcelable;
 import android.net.LinkProperties;
-import android.net.NetworkStackClient;
+import android.net.networkstack.ModuleNetworkStackClient;
 import android.os.ConditionVariable;
 
 import java.io.FileDescriptor;
@@ -75,11 +75,11 @@ public class IpClientUtil {
      *
      * <p>This is a convenience method to allow clients to use {@link IpClientCallbacks} instead of
      * {@link IIpClientCallbacks}.
-     * @see {@link NetworkStackClient#makeIpClient(String, IIpClientCallbacks)}
+     * @see {@link ModuleNetworkStackClient#makeIpClient(String, IIpClientCallbacks)}
      */
     public static void makeIpClient(Context context, String ifName, IpClientCallbacks callback) {
-        // TODO: migrate clients and remove context argument
-        NetworkStackClient.getInstance().makeIpClient(ifName, new IpClientCallbacksProxy(callback));
+        ModuleNetworkStackClient.getInstance(context)
+                .makeIpClient(ifName, new IpClientCallbacksProxy(callback));
     }
 
     /**