Merge pull request #1730 from JetBrains/feature/debug-dialog

Improve "Attach to Unity Process" dialog

rider/src/main/kotlin/com/jetbrains/rider/UnityProjectDiscoverer.kt

@@ -5,6 +5,7 @@ import com.intellij.openapi.project.Project
 import com.jetbrains.rd.platform.util.idea.LifetimedProjectService
 import com.jetbrains.rider.model.RdExistingSolution
 import com.jetbrains.rider.model.rdUnityModel
+import com.jetbrains.rider.projectView.hasSolution
 import com.jetbrains.rider.projectView.solution
 import com.jetbrains.rider.projectView.solutionDescription
 import com.jetbrains.rider.projectView.solutionFile

rider/src/main/kotlin/com/jetbrains/rider/plugins/unity/run/UnityDebuggableProcessListener.kt

@@ -0,0 +1,20 @@
+package com.jetbrains.rider.plugins.unity.run
+
+import com.intellij.openapi.project.Project
+import com.jetbrains.rd.util.lifetime.Lifetime
+import com.jetbrains.rd.util.lifetime.onTermination
+
+class UnityDebuggableProcessListener(project: Project, lifetime: Lifetime,
+                                     onProcessAdded: (UnityProcess) -> Unit,
+                                     onProcessRemoved: (UnityProcess) -> Unit) {
+
+    private val editorListener: UnityEditorListener = UnityEditorListener(project, onProcessAdded, onProcessRemoved)
+    private val playerListener: UnityPlayerListener = UnityPlayerListener(onProcessAdded, onProcessRemoved)
+
+    init {
+        lifetime.onTermination {
+            editorListener.stop()
+            playerListener.stop()
+        }
+    }
+}

rider/src/main/kotlin/com/jetbrains/rider/plugins/unity/run/UnityEditorListener.kt

@@ -0,0 +1,71 @@
+package com.jetbrains.rider.plugins.unity.run
+
+import com.intellij.execution.process.OSProcessUtil
+import com.intellij.openapi.diagnostic.Logger
+import com.intellij.openapi.project.Project
+import java.util.*
+
+class UnityEditorListener(private val project: Project,
+                          private val onEditorAdded: (UnityProcess) -> Unit,
+                          private val onEditorRemoved: (UnityProcess) -> Unit) {
+
+    companion object {
+        private val logger = Logger.getInstance(UnityEditorListener::class.java)
+    }
+
+    // Refresh once a second
+    private val refreshPeriod: Long = 1000
+    private val editorProcesses = mutableMapOf<Int, UnityLocalProcess>()
+
+    private val refreshTimer: Timer
+
+    init {
+        refreshTimer = kotlin.concurrent.timer("Listen for Unity Editor processes", true, 0L, refreshPeriod) {
+            refreshUnityEditorProcesses()
+        }
+    }
+
+    fun stop() {
+        refreshTimer.cancel()
+    }
+
+    private fun refreshUnityEditorProcesses() {
+        val start = System.currentTimeMillis()
+        logger.trace("Refreshing local editor processes...")
+
+        val processes = OSProcessUtil.getProcessList().filter {
+            UnityRunUtil.isUnityEditorProcess(it)
+        }
+
+        editorProcesses.keys.filterNot { existingEditorPid ->
+            processes.any { p -> p.pid == existingEditorPid }
+        }.forEach { p ->
+            editorProcesses[p]?.let {
+                logger.trace("Removing old Unity editor ${it.displayName}:${it.pid}")
+                onEditorRemoved(it)
+            }
+        }
+
+        val newProcesses = processes.filter { !editorProcesses.containsKey(it.pid) }
+        val unityProcessInfoMap = UnityRunUtil.getAllUnityProcessInfo(newProcesses, project)
+        newProcesses.forEach { processInfo ->
+            val unityProcessInfo = unityProcessInfoMap[processInfo.pid]
+            val editorProcess = if (!unityProcessInfo?.roleName.isNullOrEmpty()) {
+                UnityEditorHelper(processInfo.executableName, unityProcessInfo?.roleName!!, processInfo.pid, unityProcessInfo.projectName)
+            }
+            else {
+                UnityEditor(processInfo.executableName, processInfo.pid, unityProcessInfo?.projectName)
+            }
+
+            logger.trace("Adding Unity editor process ${editorProcess.displayName}:${editorProcess.pid}")
+
+            editorProcesses[processInfo.pid] = editorProcess
+            onEditorAdded(editorProcess)
+        }
+
+        if (logger.isTraceEnabled) {
+            val duration = System.currentTimeMillis() - start
+            logger.trace("Finished refreshing local editor processes. Took ${duration}ms")
+        }
+    }
+}

rider/src/main/kotlin/com/jetbrains/rider/plugins/unity/run/UnityPlayerListener.kt

@@ -1,18 +1,17 @@
 package com.jetbrains.rider.plugins.unity.run
 
-import com.intellij.execution.process.OSProcessUtil
 import com.intellij.openapi.diagnostic.Logger
-import com.intellij.openapi.project.Project
-import com.jetbrains.rd.util.lifetime.Lifetime
-import com.jetbrains.rd.util.lifetime.onTermination
 import com.jetbrains.rider.plugins.unity.util.convertPidToDebuggerPort
 import java.net.*
+import java.nio.ByteBuffer
+import java.nio.channels.DatagramChannel
+import java.nio.channels.SelectionKey
+import java.nio.channels.Selector
 import java.util.*
 import java.util.regex.Pattern
 
-class UnityPlayerListener(private val project: Project,
-                          private val onPlayerAdded: (UnityPlayer) -> Unit,
-                          private val onPlayerRemoved: (UnityPlayer) -> Unit, lifetime: Lifetime) {
+class UnityPlayerListener(private val onPlayerAdded: (UnityProcess) -> Unit,
+                          private val onPlayerRemoved: (UnityProcess) -> Unit) {
 
     companion object {
         private val logger = Logger.getInstance(UnityPlayerListener::class.java)
@@ -54,147 +53,163 @@ class UnityPlayerListener(private val project: Project,
 """, Pattern.COMMENTS)
     }
 
-    private val defaultHeartbeat = 30
+    // Refresh once a second. If a player hasn't been seen for 3 iterations, remove it from the list
+    private val refreshPeriod: Long = 1000
+    private val defaultHeartbeat = 3
 
-    private val refreshPeriod: Long = 100
-    private val waitOnSocketLength = 50
     private val multicastPorts = listOf(54997, 34997, 57997, 58997)
-    private val playerMulticastGroup = "225.0.0.222"
-    private val multicastSockets = mutableListOf<MulticastSocket>()
+    private val playerMulticastGroup = InetAddress.getByName("225.0.0.222")
+    private val selector = Selector.open()
 
     private val unityPlayerDescriptorsHeartbeats = mutableMapOf<String, Int>()
-    private val unityPlayers = mutableMapOf<String, UnityPlayer>()
+    private val unityProcesses = mutableMapOf<String, UnityProcess>()
 
     private val refreshTimer: Timer
-    private val socketsLock = Object()
+    private val syncLock = Object()
 
     init {
+
         for (networkInterface in NetworkInterface.getNetworkInterfaces()) {
             if (!networkInterface.isUp || !networkInterface.supportsMulticast()
                     || !networkInterface.inetAddresses.asSequence().any { it is Inet4Address }) {
                 continue
             }
 
-            synchronized(socketsLock) {
-                for (port in multicastPorts) {
-                    try {
-                        val multicastSocket = MulticastSocket(port)
-                        val address = InetAddress.getByName(playerMulticastGroup)
-                        multicastSocket.reuseAddress = true
-                        multicastSocket.networkInterface = networkInterface
-                        multicastSocket.soTimeout = waitOnSocketLength
-                        multicastSocket.joinGroup(address)
-                        multicastSockets.add(multicastSocket)
-                    } catch (e: Exception) {
-                        logger.warn(e.message)
-                    }
+            multicastPorts.forEach { port ->
+                try {
+                    // Setting the network interface will set the first IPv4 address on the socket's fd
+                    val channel = DatagramChannel.open(StandardProtocolFamily.INET)
+                        .setOption(StandardSocketOptions.SO_REUSEADDR, true)
+                        .setOption(StandardSocketOptions.IP_MULTICAST_IF, networkInterface)
+                        .bind(InetSocketAddress(port))
+                    channel.configureBlocking(false)
+                    channel.join(playerMulticastGroup, networkInterface)
+                    channel.register(selector, SelectionKey.OP_READ)
+                } catch (e: Exception) {
+                    logger.warn(e.message)
                 }
             }
         }
 
-        addLocalProcesses()
-        refreshTimer = startAddingPlayers()
-
-        lifetime.onTermination {
-            close()
+        refreshTimer = kotlin.concurrent.timer("Listen for Unity Players", true, 0L, refreshPeriod) {
+            refreshUnityPlayersList()
         }
     }
 
-    private fun addLocalProcesses() {
-        val unityProcesses = OSProcessUtil.getProcessList().filter { UnityRunUtil.isUnityEditorProcess(it) }
-        val unityProcessInfoMap = UnityRunUtil.getAllUnityProcessInfo(unityProcesses, project)
-        unityProcesses.map { processInfo ->
-            val unityProcessInfo = unityProcessInfoMap[processInfo.pid]
-            val port = convertPidToDebuggerPort(processInfo.pid)
-            UnityPlayer.createEditorPlayer("127.0.0.1", port, processInfo.executableName, processInfo.pid,
-                unityProcessInfo?.projectName, unityProcessInfo?.roleName)
-        }.forEach {
-            onPlayerAdded(it)
-        }
-    }
+    fun stop() {
+        refreshTimer.cancel()
 
-    private fun startAddingPlayers(): Timer {
-        return kotlin.concurrent.timer("Listen for Unity Players", true, 0L, refreshPeriod) {
-            refreshUnityPlayersList()
+        synchronized(syncLock) {
+            selector.keys().forEach {
+                try {
+                    // Close the channel. This will cancel the selection key and removes multicast group membership. It
+                    // doesn't close the socket, as there are still selector registrations active
+                    it.channel().close()
+                } catch (e: Throwable) {
+                    logger.warn(e)
+                }
+            }
+
+            // Close the selector. This deregisters the selector from all channels, and then kills the socket attached to
+            // the already closed channel
+            selector.close()
         }
     }
 
-    private fun parseUnityPlayer(unityPlayerDescriptor: String, hostAddress: String): UnityPlayer? {
+    private fun parseUnityPlayer(unityPlayerDescriptor: String, hostAddress: InetAddress): UnityProcess? {
         try {
             val matcher = unityPlayerDescriptorRegex.matcher(unityPlayerDescriptor)
             if (matcher.find()) {
-//                val ip = matcher.group("ip")
-                val port = matcher.group("port").toInt()
-                val flags = matcher.group("flags").toLong()
-                val guid = matcher.group("guid").toLong()
-                val editorGuid = matcher.group("editorId").toLong()
-                val version = matcher.group("version").toInt()
                 val id = matcher.group("id")
                 val allowDebugging = matcher.group("debug").startsWith("1")
-                // Guid's not actually a pid, but it's what we have to do
+                val guid = matcher.group("guid").toLong()
                 val debuggerPort = matcher.group("debuggerPort")?.toIntOrNull() ?: convertPidToDebuggerPort(guid)
-                val packageName: String? = matcher.group("packageName")
                 val projectName: String? = matcher.group("projectName")
 
-                // We use hostAddress instead of ip because this is the address we actually received the mulitcast from.
+                // We use hostAddress instead of ip because this is the address we actually received the multicast from.
                 // This is more accurate than what we're told, because the Unity process might be reporting the IP
                 // address of an interface that isn't reachable. For example, the iPhone player can report the local IP
                 // address of the mobile data network, which we can't reach from the current network (if we disable
                 // mobile data it works as expected)
-                return UnityPlayer(hostAddress, port, debuggerPort, flags, guid, editorGuid, version, id,
-                    allowDebugging, packageName, projectName)
+                return if (isLocalAddress(hostAddress)) {
+                    UnityLocalPlayer(id, hostAddress.hostAddress, debuggerPort, allowDebugging, projectName)
+                }
+                else {
+                    UnityRemotePlayer(id, hostAddress.hostAddress, debuggerPort, allowDebugging, projectName)
+                }
             }
         } catch (e: Exception) {
-            logger.warn("Failed to parse Unity Player: ${e.message}")
+            logger.warn(e)
         }
         return null
     }
 
+    // https://stackoverflow.com/questions/2406341/how-to-check-if-an-ip-address-is-the-local-host-on-a-multi-homed-system
+    private fun isLocalAddress(addr: InetAddress): Boolean {
+        // Check if the address is a valid special local or loop back
+        return if (addr.isAnyLocalAddress || addr.isLoopbackAddress) true else try {
+            // Check if the address is defined on any interface
+            NetworkInterface.getByInetAddress(addr) != null
+        } catch (e: SocketException) {
+            false
+        }
+    }
+
     private fun refreshUnityPlayersList() {
-        synchronized(unityPlayerDescriptorsHeartbeats) {
+        synchronized(syncLock) {
+
+            val start = System.currentTimeMillis()
             logger.trace("Refreshing Unity players list...")
+
             for (playerDescriptor in unityPlayerDescriptorsHeartbeats.keys) {
                 val currentPlayerTimeout = unityPlayerDescriptorsHeartbeats[playerDescriptor] ?: continue
-                if (currentPlayerTimeout <= 0) {
-                    unityPlayerDescriptorsHeartbeats.remove(playerDescriptor)
-                    logger.trace("Removing old Unity player $playerDescriptor")
-                    unityPlayers.remove(playerDescriptor)?.let { onPlayerRemoved(it) }
-                } else
-                    unityPlayerDescriptorsHeartbeats[playerDescriptor] = currentPlayerTimeout - 1
+                unityPlayerDescriptorsHeartbeats[playerDescriptor] = currentPlayerTimeout - 1
             }
 
-            synchronized(socketsLock) {
-                for (socket in multicastSockets) {
-                    val buf = ByteArray(1024)
-                    val recv = DatagramPacket(buf, buf.size)
+            unityPlayerDescriptorsHeartbeats.filterValues { it <= 0 }.keys.forEach { playerDescription ->
+                unityPlayerDescriptorsHeartbeats.remove(playerDescription)
+                logger.trace("Removing old Unity player $playerDescription")
+                unityProcesses.remove(playerDescription)?.let { onPlayerRemoved(it) }
+            }
+
+            // Read all data from all channels that is currently available. There might be more than one message ready
+            // on a channel as players continuously send multicast messages, so make sure we read them all. If there is
+            // nothing available, we return immediately and start sleeping.
+            val buffer = ByteBuffer.allocate(1024)
+            while (true) {
+                val readyChannels = selector.selectNow { key ->
                     try {
-                        if (socket.isClosed)
-                            continue
-                        socket.receive(recv)
-                        val descriptor = String(buf, 0, recv.length - 1)
-                        val hostAddress = recv.address.hostAddress
-                        logger.trace("Get heartbeat on port ${socket.port} from $hostAddress: $descriptor")
+                        buffer.clear()
+
+                        val channel = key.channel() as DatagramChannel
+                        val hostAddress = channel.receive(buffer) as InetSocketAddress
+                        val descriptor = String(buffer.array(), 0, buffer.position() - 1)
+
+                        if (logger.isTraceEnabled) {
+                            logger.trace("Got heartbeat on ${channel.remoteAddress} from $hostAddress: $descriptor")
+                        }
+
                         if (!unityPlayerDescriptorsHeartbeats.containsKey(descriptor)) {
-                            parseUnityPlayer(descriptor, hostAddress)?.let {
-                                unityPlayers[descriptor] = it
-                                onPlayerAdded(it)
+                            parseUnityPlayer(descriptor, hostAddress.address)?.let { player ->
+                                unityProcesses[descriptor] = player
+                                onPlayerAdded(player)
                             }
                         }
+
                         unityPlayerDescriptorsHeartbeats[descriptor] = defaultHeartbeat
-                    } catch (e: SocketTimeoutException) {
-                        //wait timeout, go to the next port
+                    } catch (e: Exception) {
+                        logger.warn(e)
                     }
                 }
+
+                if (readyChannels == 0) {
+                    break
+                }
             }
-            logger.trace("Finished refreshing of Unity players list.")
-        }
-    }
 
-    private fun close() {
-        refreshTimer.cancel()
-        synchronized(socketsLock) {
-            for (socket in multicastSockets) {
-                socket.close()
+            if (logger.isTraceEnabled) {
+                val duration = System.currentTimeMillis() - start
+                logger.trace("Finished refreshing Unity players list. Took ${duration}ms")
             }
         }
     }