3 files changed, 32 insertions, 6 deletions
diff --git a/src/android/net/ip/IpClient.java b/src/android/net/ip/IpClient.java
index 612ebf7..4315d34 100644
--- a/src/android/net/ip/IpClient.java
+++ b/src/android/net/ip/IpClient.java
@@ -543,7 +543,7 @@ public class IpClient extends StateMachine {
     }
 
     private void startStateMachineUpdaters() {
-        mObserverRegistry.registerObserver(mLinkObserver, getHandler());
+        mObserverRegistry.registerObserverForNonblockingCallback(mLinkObserver);
     }
 
     private void stopStateMachineUpdaters() {
diff --git a/src/com/android/server/NetworkObserverRegistry.java b/src/com/android/server/NetworkObserverRegistry.java
index 86e5e14..4f55779 100644
--- a/src/com/android/server/NetworkObserverRegistry.java
+++ b/src/com/android/server/NetworkObserverRegistry.java
@@ -24,8 +24,10 @@ import android.net.LinkAddress;
 import android.net.RouteInfo;
 import android.os.Handler;
 import android.os.RemoteException;
+import android.util.Log;
 
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
@@ -35,6 +37,7 @@ import java.util.concurrent.ConcurrentHashMap;
  * all INetworkManagementEventObserver objects that have registered with it.
  */
 public class NetworkObserverRegistry extends INetdUnsolicitedEventListener.Stub {
+    private static final String TAG = NetworkObserverRegistry.class.getSimpleName();
 
     /**
      * Constructs a new NetworkObserverRegistry.
@@ -53,7 +56,7 @@ public class NetworkObserverRegistry extends INetdUnsolicitedEventListener.Stub
         netd.registerUnsolicitedEventListener(this);
     }
 
-    private final ConcurrentHashMap<NetworkObserver, Handler> mObservers =
+    private final ConcurrentHashMap<NetworkObserver, Optional<Handler>> mObservers =
             new ConcurrentHashMap<>();
 
     /**
@@ -61,7 +64,20 @@ public class NetworkObserverRegistry extends INetdUnsolicitedEventListener.Stub
      * This method may be called on any thread.
      */
     public void registerObserver(@NonNull NetworkObserver observer, @NonNull Handler handler) {
-        mObservers.put(observer, handler);
+        if (handler == null) {
+            throw new IllegalArgumentException("handler must be non-null");
+        }
+        mObservers.put(observer, Optional.of(handler));
+    }
+
+    /**
+     * Registers the specified observer, and start sending callbacks to it.
+     *
+     * <p>This method must only be called with callbacks that are nonblocking, such as callbacks
+     * that only send a message to a StateMachine.
+     */
+    public void registerObserverForNonblockingCallback(@NonNull NetworkObserver observer) {
+        mObservers.put(observer, Optional.empty());
     }
 
     /**
@@ -80,9 +96,19 @@ public class NetworkObserverRegistry extends INetdUnsolicitedEventListener.Stub
     private void invokeForAllObservers(@NonNull final NetworkObserverEventCallback callback) {
         // ConcurrentHashMap#entrySet is weakly consistent: observers that were in the map before
         // creation will be processed, those added during traversal may or may not.
-        for (Map.Entry<NetworkObserver, Handler> entry : mObservers.entrySet()) {
+        for (Map.Entry<NetworkObserver, Optional<Handler>> entry : mObservers.entrySet()) {
             final NetworkObserver observer = entry.getKey();
-            entry.getValue().post(() -> callback.sendCallback(observer));
+            final Optional<Handler> handler = entry.getValue();
+            if (handler.isPresent()) {
+                handler.get().post(() -> callback.sendCallback(observer));
+                return;
+            }
+
+            try {
+                callback.sendCallback(observer);
+            } catch (RuntimeException e) {
+                Log.e(TAG, "Error sending callback to observer", e);
+            }
         }
     }
 
diff --git a/tests/src/android/net/ip/IpClientTest.java b/tests/src/android/net/ip/IpClientTest.java
index 40d510a..7e57d1e 100644
--- a/tests/src/android/net/ip/IpClientTest.java
+++ b/tests/src/android/net/ip/IpClientTest.java
@@ -129,7 +129,7 @@ public class IpClientTest {
         verify(mNetd, timeout(TEST_TIMEOUT_MS).times(1)).interfaceSetEnableIPv6(ifname, false);
         verify(mNetd, timeout(TEST_TIMEOUT_MS).times(1)).interfaceClearAddrs(ifname);
         ArgumentCaptor<NetworkObserver> arg = ArgumentCaptor.forClass(NetworkObserver.class);
-        verify(mObserverRegistry, times(1)).registerObserver(arg.capture(), any());
+        verify(mObserverRegistry, times(1)).registerObserverForNonblockingCallback(arg.capture());
         mObserver = arg.getValue();
         reset(mObserverRegistry);
         reset(mNetd);