diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index f1c68cb..61c4fb1 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -5577,14 +5577,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: {
@@ -9865,9 +9865,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()