diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 4d03f666f2..96197cea93 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -7883,12 +7883,14 @@ public class ConnectivityService extends IConnectivityManager.Stub
         ensureRunningOnConnectivityServiceThread();
         final IBinder iCb = cb.asBinder();
 
-        if (!mConnectivityDiagnosticsCallbacks.containsKey(iCb)) {
+        final ConnectivityDiagnosticsCallbackInfo cbInfo =
+                mConnectivityDiagnosticsCallbacks.remove(iCb);
+        if (cbInfo == null) {
             if (VDBG) log("Removing diagnostics callback that is not currently registered");
             return;
         }
 
-        final NetworkRequestInfo nri = mConnectivityDiagnosticsCallbacks.get(iCb).mRequestInfo;
+        final NetworkRequestInfo nri = cbInfo.mRequestInfo;
 
         if (uid != nri.mUid) {
             if (VDBG) loge("Different uid than registrant attempting to unregister cb");
@@ -7900,8 +7902,6 @@ public class ConnectivityService extends IConnectivityManager.Stub
         // enforceRequestCountLimit().
         decrementNetworkRequestPerUidCount(nri);
 
-        final ConnectivityDiagnosticsCallbackInfo cbInfo =
-                mConnectivityDiagnosticsCallbacks.remove(iCb);
         iCb.unlinkToDeath(cbInfo, 0);
     }