diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 916a68d398..0514532e45 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -7853,12 +7853,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");
@@ -7870,8 +7872,6 @@ public class ConnectivityService extends IConnectivityManager.Stub
         // enforceRequestCountLimit().
         decrementNetworkRequestPerUidCount(nri);
 
-        final ConnectivityDiagnosticsCallbackInfo cbInfo =
-                mConnectivityDiagnosticsCallbacks.remove(iCb);
         iCb.unlinkToDeath(cbInfo, 0);
     }