diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 4224da9..17389b4 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -2279,23 +2279,12 @@
         private final ISocketKeepaliveCallback mCallback;
         private final ExecutorService mExecutor;
 
-        private volatile Integer mSlot;
-
         @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
         public void stop() {
             try {
                 mExecutor.execute(() -> {
                     try {
-                        if (mSlot != null) {
-                            // TODO : this is incorrect, because in the presence of auto on/off
-                            // keepalive the slot associated with this keepalive can have
-                            // changed. Also, this actually causes another problem where some other
-                            // app might stop your keepalive if it just knows the network and
-                            // the slot and goes through the trouble of grabbing the aidl object.
-                            // This code should use the callback to identify what keepalive to
-                            // stop instead.
-                            mService.stopKeepalive(mNetwork, mSlot);
-                        }
+                        mService.stopKeepalive(mCallback);
                     } catch (RemoteException e) {
                         Log.e(TAG, "Error stopping packet keepalive: ", e);
                         throw e.rethrowFromSystemServer();
@@ -2313,11 +2302,10 @@
             mExecutor = Executors.newSingleThreadExecutor();
             mCallback = new ISocketKeepaliveCallback.Stub() {
                 @Override
-                public void onStarted(int slot) {
+                public void onStarted() {
                     final long token = Binder.clearCallingIdentity();
                     try {
                         mExecutor.execute(() -> {
-                            mSlot = slot;
                             callback.onStarted();
                         });
                     } finally {
@@ -2330,7 +2318,6 @@
                     final long token = Binder.clearCallingIdentity();
                     try {
                         mExecutor.execute(() -> {
-                            mSlot = null;
                             callback.onStopped();
                         });
                     } finally {
@@ -2344,7 +2331,6 @@
                     final long token = Binder.clearCallingIdentity();
                     try {
                         mExecutor.execute(() -> {
-                            mSlot = null;
                             callback.onError(error);
                         });
                     } finally {
diff --git a/framework/src/android/net/IConnectivityManager.aidl b/framework/src/android/net/IConnectivityManager.aidl
index 7db231e..acbc31e 100644
--- a/framework/src/android/net/IConnectivityManager.aidl
+++ b/framework/src/android/net/IConnectivityManager.aidl
@@ -193,7 +193,7 @@
     void startTcpKeepalive(in Network network, in ParcelFileDescriptor pfd, int intervalSeconds,
             in ISocketKeepaliveCallback cb);
 
-    void stopKeepalive(in Network network, int slot);
+    void stopKeepalive(in ISocketKeepaliveCallback cb);
 
     String getCaptivePortalServerUrl();
 
diff --git a/framework/src/android/net/ISocketKeepaliveCallback.aidl b/framework/src/android/net/ISocketKeepaliveCallback.aidl
index 020fbca..1240e37 100644
--- a/framework/src/android/net/ISocketKeepaliveCallback.aidl
+++ b/framework/src/android/net/ISocketKeepaliveCallback.aidl
@@ -24,7 +24,7 @@
 oneway interface ISocketKeepaliveCallback
 {
     /** The keepalive was successfully started. */
-    void onStarted(int slot);
+    void onStarted();
     /** The keepalive was successfully stopped. */
     void onStopped();
     /** The keepalive was stopped because of an error. */
diff --git a/framework/src/android/net/NattSocketKeepalive.java b/framework/src/android/net/NattSocketKeepalive.java
index 4d45e70..77137f4 100644
--- a/framework/src/android/net/NattSocketKeepalive.java
+++ b/framework/src/android/net/NattSocketKeepalive.java
@@ -91,9 +91,7 @@
     protected void stopImpl() {
         mExecutor.execute(() -> {
             try {
-                if (mSlot != null) {
-                    mService.stopKeepalive(mNetwork, mSlot);
-                }
+                mService.stopKeepalive(mCallback);
             } catch (RemoteException e) {
                 Log.e(TAG, "Error stopping socket keepalive: ", e);
                 throw e.rethrowFromSystemServer();
diff --git a/framework/src/android/net/NetworkAgent.java b/framework/src/android/net/NetworkAgent.java
index 3ec00d9..8fe20de 100644
--- a/framework/src/android/net/NetworkAgent.java
+++ b/framework/src/android/net/NetworkAgent.java
@@ -291,7 +291,9 @@
     /**
      * Requests that the specified keepalive packet be stopped.
      *
-     * arg1 = hardware slot number of the keepalive to stop.
+     * arg1 = unused
+     * arg2 = error code (SUCCESS)
+     * obj = callback to identify the keepalive
      *
      * Also used internally by ConnectivityService / KeepaliveTracker, with different semantics.
      * @hide
diff --git a/framework/src/android/net/SocketKeepalive.java b/framework/src/android/net/SocketKeepalive.java
index 90e5e9b..2911ce7 100644
--- a/framework/src/android/net/SocketKeepalive.java
+++ b/framework/src/android/net/SocketKeepalive.java
@@ -21,7 +21,6 @@
 import android.annotation.IntDef;
 import android.annotation.IntRange;
 import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.os.Binder;
 import android.os.ParcelFileDescriptor;
@@ -249,9 +248,6 @@
     @NonNull protected final Executor mExecutor;
     /** @hide */
     @NonNull protected final ISocketKeepaliveCallback mCallback;
-    // TODO: remove slot since mCallback could be used to identify which keepalive to stop.
-    /** @hide */
-    @Nullable protected Integer mSlot;
 
     /** @hide */
     public SocketKeepalive(@NonNull IConnectivityManager service, @NonNull Network network,
@@ -263,11 +259,10 @@
         mExecutor = executor;
         mCallback = new ISocketKeepaliveCallback.Stub() {
             @Override
-            public void onStarted(int slot) {
+            public void onStarted() {
                 final long token = Binder.clearCallingIdentity();
                 try {
                     mExecutor.execute(() -> {
-                        mSlot = slot;
                         callback.onStarted();
                     });
                 } finally {
@@ -280,7 +275,6 @@
                 final long token = Binder.clearCallingIdentity();
                 try {
                     executor.execute(() -> {
-                        mSlot = null;
                         callback.onStopped();
                     });
                 } finally {
@@ -293,7 +287,6 @@
                 final long token = Binder.clearCallingIdentity();
                 try {
                     executor.execute(() -> {
-                        mSlot = null;
                         callback.onError(error);
                     });
                 } finally {
@@ -306,7 +299,6 @@
                 final long token = Binder.clearCallingIdentity();
                 try {
                     executor.execute(() -> {
-                        mSlot = null;
                         callback.onDataReceived();
                     });
                 } finally {
diff --git a/framework/src/android/net/TcpSocketKeepalive.java b/framework/src/android/net/TcpSocketKeepalive.java
index 51d805e..cda5830 100644
--- a/framework/src/android/net/TcpSocketKeepalive.java
+++ b/framework/src/android/net/TcpSocketKeepalive.java
@@ -69,9 +69,7 @@
     protected void stopImpl() {
         mExecutor.execute(() -> {
             try {
-                if (mSlot != null) {
-                    mService.stopKeepalive(mNetwork, mSlot);
-                }
+                mService.stopKeepalive(mCallback);
             } catch (RemoteException e) {
                 Log.e(TAG, "Error stopping packet keepalive: ", e);
                 throw e.rethrowFromSystemServer();
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 2371911..a570ab1 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -5580,14 +5580,14 @@
                 }
                 // Sent by KeepaliveTracker to process an app request on the state machine thread.
                 case NetworkAgent.CMD_STOP_SOCKET_KEEPALIVE: {
-                    NetworkAgentInfo nai = getNetworkAgentInfoForNetwork((Network) msg.obj);
-                    if (nai == null) {
-                        Log.e(TAG, "Attempt to stop keepalive on nonexistent network");
+                    final AutomaticOnOffKeepalive ki = mKeepaliveTracker.getKeepaliveForBinder(
+                            (IBinder) msg.obj);
+                    if (ki == null) {
+                        Log.e(TAG, "Attempt to stop an already stopped keepalive");
                         return;
                     }
-                    int slot = msg.arg1;
-                    int reason = msg.arg2;
-                    mKeepaliveTracker.handleStopKeepalive(nai, slot, reason);
+                    final int reason = msg.arg2;
+                    mKeepaliveTracker.handleStopKeepalive(ki, reason);
                     break;
                 }
                 case EVENT_REPORT_NETWORK_CONNECTIVITY: {
@@ -9870,9 +9870,10 @@
     }
 
     @Override
-    public void stopKeepalive(Network network, int slot) {
+    public void stopKeepalive(@NonNull final ISocketKeepaliveCallback cb) {
         mHandler.sendMessage(mHandler.obtainMessage(
-                NetworkAgent.CMD_STOP_SOCKET_KEEPALIVE, slot, SocketKeepalive.SUCCESS, network));
+                NetworkAgent.CMD_STOP_SOCKET_KEEPALIVE, 0, SocketKeepalive.SUCCESS,
+                Objects.requireNonNull(cb).asBinder()));
     }
 
     @Override
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index 18e2dd8..9dcbac8 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -329,30 +329,6 @@
         handleResumeKeepalive(newKi);
     }
 
-    // TODO : this method should be removed ; the keepalives should always be indexed by callback
-    private int findAutomaticOnOffKeepaliveIndex(@NonNull Network network, int slot) {
-        ensureRunningOnHandlerThread();
-
-        int index = 0;
-        for (AutomaticOnOffKeepalive ki : mAutomaticOnOffKeepalives) {
-            if (ki.match(network, slot)) {
-                return index;
-            }
-            index++;
-        }
-        return -1;
-    }
-
-    // TODO : this method should be removed ; the keepalives should always be indexed by callback
-    @Nullable
-    private AutomaticOnOffKeepalive findAutomaticOnOffKeepalive(@NonNull Network network,
-            int slot) {
-        ensureRunningOnHandlerThread();
-
-        final int index = findAutomaticOnOffKeepaliveIndex(network, slot);
-        return (index >= 0) ? mAutomaticOnOffKeepalives.get(index) : null;
-    }
-
     /**
      * Find the AutomaticOnOffKeepalive associated with a given callback.
      * @return the keepalive associated with this callback, or null if none
@@ -415,17 +391,12 @@
     /**
      * Handle stop keepalives on the specific network with given slot.
      */
-    public void handleStopKeepalive(@NonNull NetworkAgentInfo nai, int slot, int reason) {
-        final AutomaticOnOffKeepalive autoKi = findAutomaticOnOffKeepalive(nai.network, slot);
-        if (null == autoKi) {
-            Log.e(TAG, "Attempt to stop nonexistent keepalive " + slot + " on " + nai);
-            return;
-        }
-
+    public void handleStopKeepalive(@NonNull final AutomaticOnOffKeepalive autoKi, int reason) {
         // Stop the keepalive unless it was suspended. This includes the case where it's managed
         // but enabled, and the case where it's always on.
         if (autoKi.mAutomaticOnOffState != STATE_SUSPENDED) {
-            mKeepaliveTracker.handleStopKeepalive(nai, slot, reason);
+            final KeepaliveTracker.KeepaliveInfo ki = autoKi.mKi;
+            mKeepaliveTracker.handleStopKeepalive(ki.getNai(), ki.getSlot(), reason);
         }
 
         cleanupAutoOnOffKeepalive(autoKi);
diff --git a/service/src/com/android/server/connectivity/KeepaliveTracker.java b/service/src/com/android/server/connectivity/KeepaliveTracker.java
index 7cb613b..a512b7c 100644
--- a/service/src/com/android/server/connectivity/KeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveTracker.java
@@ -589,9 +589,9 @@
                 Log.d(TAG, "Started keepalive " + slot + " on " + nai.toShortString());
                 ki.mStartedState = KeepaliveInfo.STARTED;
                 try {
-                    ki.mCallback.onStarted(slot);
+                    ki.mCallback.onStarted();
                 } catch (RemoteException e) {
-                    Log.w(TAG, "Discarded onStarted(" + slot + ") callback");
+                    Log.w(TAG, "Discarded onStarted callback");
                 }
             } else {
                 Log.d(TAG, "Failed to start keepalive " + slot + " on " + nai.toShortString()