Rewrite MdnsDiscoveryAgent to treat the singleton JmmDNS object safely

moonlight-common/src/com/limelight/nvstream/mdns/MdnsDiscoveryAgent.java

@@ -8,8 +8,6 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Timer;
-import java.util.TimerTask;
 
 import javax.jmdns.JmmDNS;
 import javax.jmdns.ServiceEvent;
@@ -18,50 +16,95 @@ import javax.jmdns.ServiceListener;
 
 import com.limelight.LimeLog;
 
-public class MdnsDiscoveryAgent {
+public class MdnsDiscoveryAgent implements ServiceListener {
 	public static final String SERVICE_TYPE = "_nvstream._tcp.local.";
 	
-	private JmmDNS resolver;
-	private HashMap<InetAddress, MdnsComputer> computers;
 	private MdnsDiscoveryListener listener;
-	private HashSet<String> pendingResolution;
+	private Thread discoveryThread;
-	private boolean stop;
+	private HashMap<InetAddress, MdnsComputer> computers = new HashMap<InetAddress, MdnsComputer>();
-	private ServiceListener nvstreamListener = new ServiceListener() {
+	private HashSet<String> pendingResolution = new HashSet<String>();
+	
+	// The resolver factory's instance member has a static lifetime which
+	// means our ref count and listener must be static also.
+	private static int resolverRefCount = 0;
+	private static HashSet<ServiceListener> listeners = new HashSet<ServiceListener>();
+	private static ServiceListener nvstreamListener = new ServiceListener() {
+		@Override
 		public void serviceAdded(ServiceEvent event) {
-			LimeLog.info("mDNS: Machine appeared: "+event.getInfo().getName());
+			HashSet<ServiceListener> localListeners;
 			
-			ServiceInfo info = event.getDNS().getServiceInfo(SERVICE_TYPE, event.getInfo().getName(), 500);
+			// Copy the listener set into a new set so we can invoke
-			if (info == null) {
+			// the callbacks without holding the listeners monitor the
-				// This machine is pending resolution
+			// whole time.
-				pendingResolution.add(event.getInfo().getName());
+			synchronized (listeners) {
-				return;
+				localListeners = new HashSet<ServiceListener>(listeners);
 			}
 			
-			LimeLog.info("mDNS: Resolved (blocking)");
+			for (ServiceListener listener : localListeners) {
-			handleResolvedServiceInfo(info);
+				listener.serviceAdded(event);
+			}
 		}
 
+		@Override
 		public void serviceRemoved(ServiceEvent event) {
-			LimeLog.info("mDNS: Machine disappeared: "+event.getInfo().getName());
+			HashSet<ServiceListener> localListeners;
 			
-			Inet4Address addrs[] = event.getInfo().getInet4Addresses();
+			// Copy the listener set into a new set so we can invoke
-			for (Inet4Address addr : addrs) {
+			// the callbacks without holding the listeners monitor the
-				synchronized (computers) {
+			// whole time.
-					MdnsComputer computer = computers.remove(addr);
+			synchronized (listeners) {
-					if (computer != null) {
+				localListeners = new HashSet<ServiceListener>(listeners);
-						listener.notifyComputerRemoved(computer);
+			}
-						break;
+			
-					}
+			for (ServiceListener listener : localListeners) {
-				}
+				listener.serviceRemoved(event);
 			}
 		}
 
+		@Override
 		public void serviceResolved(ServiceEvent event) {
-			LimeLog.info("mDNS: Machine resolved (callback): "+event.getInfo().getName());
+			HashSet<ServiceListener> localListeners;
-			handleResolvedServiceInfo(event.getInfo());
+			
+			// Copy the listener set into a new set so we can invoke
+			// the callbacks without holding the listeners monitor the
+			// whole time.
+			synchronized (listeners) {
+				localListeners = new HashSet<ServiceListener>(listeners);
+			}
+			
+			for (ServiceListener listener : localListeners) {
+				listener.serviceResolved(event);
+			}
 		}
 	};
 	
+	private static JmmDNS referenceResolver() {
+		synchronized (MdnsDiscoveryAgent.class) {
+			JmmDNS instance = JmmDNS.Factory.getInstance();
+			if (++resolverRefCount == 1) {
+				// This will cause the listener to be invoked for known hosts immediately.
+				// JmDNS only supports one listener per service, so we have to do this here
+				// with a static listener.
+				instance.addServiceListener(SERVICE_TYPE, nvstreamListener);
+			}
+			return instance;
+		}
+	}
+	
+	private static void dereferenceResolver() {
+		synchronized (MdnsDiscoveryAgent.class) {
+			if (--resolverRefCount == 0) {
+				try {
+					JmmDNS.Factory.close();
+				} catch (IOException e) {}
+			}
+		}
+	}
+	
+	public MdnsDiscoveryAgent(MdnsDiscoveryListener listener) {
+		this.listener = listener;
+	}
+	
 	private void handleResolvedServiceInfo(ServiceInfo info) {
 		MdnsComputer computer;
 				
@@ -100,65 +143,67 @@ public class MdnsDiscoveryAgent {
 		return new MdnsComputer(name, address);
 	}
 	
-	public MdnsDiscoveryAgent(MdnsDiscoveryListener listener) {
-		computers = new HashMap<InetAddress, MdnsComputer>();
-		pendingResolution = new HashSet<String>();
-		this.listener = listener;
-	}
-	
 	public void startDiscovery(final int discoveryIntervalMs) {
-		stop = false;
+		// Kill any existing discovery before starting a new one
+		stopDiscovery();
 		
-		final Timer t = new Timer();
+		// Add our listener to the set
-		t.schedule(new TimerTask() {
+		synchronized (listeners) {
+			listeners.add(MdnsDiscoveryAgent.this);
+		}
+		
+		discoveryThread = new Thread() {
 			@Override
 			public void run() {
-				synchronized (MdnsDiscoveryAgent.this) {
+				// This may result in listener callbacks so we must register
-					// Stop if requested
+				// our listener first.
-					if (stop) {
+				JmmDNS resolver = referenceResolver();
-						// There will be no further timer invocations now
+				
-						t.cancel();
+				try {
+					while (!Thread.interrupted()) {
+						// Start an mDNS request
+						resolver.requestServiceInfo(SERVICE_TYPE, null, discoveryIntervalMs);
 						
-						if (resolver != null) {
+						// Run service resolution again for pending machines
-							resolver.removeServiceListener(SERVICE_TYPE, nvstreamListener);
+						ArrayList<String> pendingNames = new ArrayList<String>(pendingResolution);
-							try {
+						for (String name : pendingNames) {
-								JmmDNS.Factory.close();
+							LimeLog.info("mDNS: Retrying service resolution for machine: "+name);
-							} catch (IOException e) {}
+							ServiceInfo[] infos = resolver.getServiceInfos(SERVICE_TYPE, name, 500);
-							resolver = null;
+							if (infos != null && infos.length != 0) {
+								LimeLog.info("mDNS: Resolved (retry) with "+infos.length+" service entries");
+								handleResolvedServiceInfo(infos[0]);
+							}
 						}
-						return;
-					}
-					
-					// Perform initialization
-					if (resolver == null) {
-						resolver = JmmDNS.Factory.getInstance();
 						
-						// This will cause the listener to be invoked for known hosts immediately.
+						// Wait for the next polling interval
-						// We do this in the timer callback to be off the main thread when this happens.
+						try {
-						resolver.addServiceListener(SERVICE_TYPE, nvstreamListener);
+							Thread.sleep(discoveryIntervalMs);
-					}
+						} catch (InterruptedException e) {
-					
+							break;
-					// Send another mDNS query
-					resolver.requestServiceInfo(SERVICE_TYPE, null, discoveryIntervalMs);
-
-					// Run service resolution again for pending machines
-					ArrayList<String> pendingNames = new ArrayList<String>(pendingResolution);
-					for (String name : pendingNames) {
-						LimeLog.info("mDNS: Retrying service resolution for machine: "+name);
-						ServiceInfo[] infos = resolver.getServiceInfos(SERVICE_TYPE, name, 500);
-						if (infos != null && infos.length != 0) {
-							LimeLog.info("mDNS: Resolved (retry) with "+infos.length+" service entries");
-							handleResolvedServiceInfo(infos[0]);
 						}
 					}
 				}
+				finally {
+					// Dereference the resolver
+					dereferenceResolver();
+				}
 			}
-		}, 0, discoveryIntervalMs);
+		};
+		discoveryThread.setName("mDNS Discovery Thread");
+		discoveryThread.start();
 	}
 	
 	public void stopDiscovery() {
-		// Trigger a stop on the next timer expiration
+		// Remove our listener from the set
-		stop = true;
+		synchronized (listeners) {
+			listeners.remove(MdnsDiscoveryAgent.this);
+		}
+		
+		// If there's already a running thread, interrupt it
+		if (discoveryThread != null) {
+			discoveryThread.interrupt();
+			discoveryThread = null;
+		}
 	}
 	
 	public List<MdnsComputer> getComputerSet() {
@@ -166,4 +211,41 @@ public class MdnsDiscoveryAgent {
 			return new ArrayList<MdnsComputer>(computers.values());
 		}
 	}
+
+	@Override
+	public void serviceAdded(ServiceEvent event) {
+		LimeLog.info("mDNS: Machine appeared: "+event.getInfo().getName());
+
+		ServiceInfo info = event.getDNS().getServiceInfo(SERVICE_TYPE, event.getInfo().getName(), 500);
+		if (info == null) {
+			// This machine is pending resolution
+			pendingResolution.add(event.getInfo().getName());
+			return;
+		}
+		
+		LimeLog.info("mDNS: Resolved (blocking)");
+		handleResolvedServiceInfo(info);
+	}
+
+	@Override
+	public void serviceRemoved(ServiceEvent event) {
+		LimeLog.info("mDNS: Machine disappeared: "+event.getInfo().getName());
+		
+		Inet4Address addrs[] = event.getInfo().getInet4Addresses();
+		for (Inet4Address addr : addrs) {
+			synchronized (computers) {
+				MdnsComputer computer = computers.remove(addr);
+				if (computer != null) {
+					listener.notifyComputerRemoved(computer);
+					break;
+				}
+			}
+		}
+	}
+
+	@Override
+	public void serviceResolved(ServiceEvent event) {
+		LimeLog.info("mDNS: Machine resolved (callback): "+event.getInfo().getName());
+		handleResolvedServiceInfo(event.getInfo());
+	}
 }