impl: support for displaying network latency

src/main/kotlin/com/coder/toolbox/CoderRemoteEnvironment.kt

@@ -2,15 +2,18 @@ package com.coder.toolbox
 
 import com.coder.toolbox.browser.BrowserUtil
 import com.coder.toolbox.cli.CoderCLIManager
+import com.coder.toolbox.cli.SshCommandProcessHandle
 import com.coder.toolbox.models.WorkspaceAndAgentStatus
 import com.coder.toolbox.sdk.CoderRestClient
 import com.coder.toolbox.sdk.ex.APIResponseException
+import com.coder.toolbox.sdk.v2.models.NetworkMetrics
 import com.coder.toolbox.sdk.v2.models.Workspace
 import com.coder.toolbox.sdk.v2.models.WorkspaceAgent
 import com.coder.toolbox.util.waitForFalseWithTimeout
 import com.coder.toolbox.util.withPath
 import com.coder.toolbox.views.Action
 import com.coder.toolbox.views.EnvironmentView
+import com.jetbrains.toolbox.api.localization.LocalizableString
 import com.jetbrains.toolbox.api.remoteDev.AfterDisconnectHook
 import com.jetbrains.toolbox.api.remoteDev.BeforeConnectionHook
 import com.jetbrains.toolbox.api.remoteDev.DeleteEnvironmentConfirmationParams
@@ -20,15 +23,21 @@ import com.jetbrains.toolbox.api.remoteDev.environments.EnvironmentContentsView
 import com.jetbrains.toolbox.api.remoteDev.states.EnvironmentDescription
 import com.jetbrains.toolbox.api.remoteDev.states.RemoteEnvironmentState
 import com.jetbrains.toolbox.api.ui.actions.ActionDescription
+import com.squareup.moshi.Moshi
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withTimeout
+import java.io.File
+import java.nio.file.Path
 import kotlin.time.Duration.Companion.minutes
 import kotlin.time.Duration.Companion.seconds
 
+private val POLL_INTERVAL = 5.seconds
+
 /**
  * Represents an agent and workspace combination.
  *
@@ -44,17 +53,20 @@ class CoderRemoteEnvironment(
     private var wsRawStatus = WorkspaceAndAgentStatus.from(workspace, agent)
 
     override var name: String = "${workspace.name}.${agent.name}"
-
     private var isConnected: MutableStateFlow<Boolean> = MutableStateFlow(false)
     override val connectionRequest: MutableStateFlow<Boolean> = MutableStateFlow(false)
 
     override val state: MutableStateFlow<RemoteEnvironmentState> =
         MutableStateFlow(wsRawStatus.toRemoteEnvironmentState(context))
     override val description: MutableStateFlow<EnvironmentDescription> =
         MutableStateFlow(EnvironmentDescription.General(context.i18n.pnotr(workspace.templateDisplayName)))
-
+    override val additionalEnvironmentInformation: MutableMap<LocalizableString, String> = mutableMapOf()
     override val actionsList: MutableStateFlow<List<ActionDescription>> = MutableStateFlow(getAvailableActions())
 
+    private val networkMetricsMarshaller = Moshi.Builder().build().adapter(NetworkMetrics::class.java)
+    private val proxyCommandHandle = SshCommandProcessHandle(context)
+    private var pollJob: Job? = null
+
     fun asPairOfWorkspaceAndAgent(): Pair<Workspace, WorkspaceAgent> = Pair(workspace, agent)
 
     private fun getAvailableActions(): List<ActionDescription> {
@@ -141,9 +153,49 @@ class CoderRemoteEnvironment(
     override fun beforeConnection() {
         context.logger.info("Connecting to $id...")
         isConnected.update { true }
+        pollJob = pollNetworkMetrics()
+    }
+
+    private fun pollNetworkMetrics(): Job = context.cs.launch {
+        context.logger.info("Starting the network metrics poll job for $id")
+        while (isActive) {
+            context.logger.debug("Searching SSH command's PID for workspace $id...")
+            val pid = proxyCommandHandle.findByWorkspaceAndAgent(workspace, agent)
+            if (pid == null) {
+                context.logger.debug("No SSH command PID was found for workspace $id")
+                delay(POLL_INTERVAL)
+                continue
+            }
+
+            val metricsFile = Path.of(context.settingsStore.networkInfoDir, "$pid.json").toFile()
+            if (metricsFile.doesNotExists()) {
+                context.logger.debug("No metrics file found at ${metricsFile.absolutePath} for $id")
+                delay(POLL_INTERVAL)
+                continue
+            }
+            context.logger.debug("Loading metrics from ${metricsFile.absolutePath} for $id")
+            try {
+                val metrics = networkMetricsMarshaller.fromJson(metricsFile.readText())
+                if (metrics == null) {
+                    return@launch
+                }
+                context.logger.debug("$id metrics: $metrics")
+                additionalEnvironmentInformation.put(context.i18n.ptrl("Network Metrics"), metrics.toPretty())
+            } catch (e: Exception) {
+                context.logger.error(
+                    e,
+                    "Error encountered while trying to load network metrics from ${metricsFile.absolutePath} for $id"
+                )
+            }
+            delay(POLL_INTERVAL)
+        }
     }
 
+    private fun File.doesNotExists(): Boolean = !this.exists()
+
     override fun afterDisconnect() {
+        context.logger.info("Stopping the network metrics poll job for $id")
+        pollJob?.cancel()
         this.connectionRequest.update { false }
         isConnected.update { false }
         context.logger.info("Disconnected from $id")

src/main/kotlin/com/coder/toolbox/cli/CoderCLIManager.kt

@@ -271,14 +271,13 @@ class CoderCLIManager(
                 "ssh",
                 "--stdio",
                 if (settings.disableAutostart && feats.disableAutostart) "--disable-autostart" else null,
+                "--network-info-dir ${escape(settings.networkInfoDir)}"
             )
         val proxyArgs = baseArgs + listOfNotNull(
             if (!settings.sshLogDirectory.isNullOrBlank()) "--log-dir" else null,
             if (!settings.sshLogDirectory.isNullOrBlank()) escape(settings.sshLogDirectory!!) else null,
             if (feats.reportWorkspaceUsage) "--usage-app=jetbrains" else null,
         )
-        val backgroundProxyArgs =
-            baseArgs + listOfNotNull(if (feats.reportWorkspaceUsage) "--usage-app=disable" else null)
         val extraConfig =
             if (!settings.sshConfigOptions.isNullOrBlank()) {
                 "\n" + settings.sshConfigOptions!!.prependIndent("  ")

src/main/kotlin/com/coder/toolbox/cli/SshCommandProcessHandle.kt

@@ -0,0 +1,42 @@
+package com.coder.toolbox.cli
+
+import com.coder.toolbox.CoderToolboxContext
+import com.coder.toolbox.sdk.v2.models.Workspace
+import com.coder.toolbox.sdk.v2.models.WorkspaceAgent
+import kotlin.jvm.optionals.getOrNull
+
+/**
+ * Identifies the PID for the SSH Coder command spawned by Toolbox.
+ */
+class SshCommandProcessHandle(private val ctx: CoderToolboxContext) {
+
+    /**
+     * Finds the PID of a Coder (not the proxy command) ssh cmd associated with the specified workspace and agent.
+     * Null is returned when no ssh command process was found.
+     *
+     * Implementation Notes:
+     * An iterative DFS approach where we start with Toolbox's direct children, grep the command
+     * and if nothing is found we continue with the processes children. Toolbox spawns an ssh command
+     * as a separate command which in turns spawns another child for the proxy command.
+     */
+    fun findByWorkspaceAndAgent(ws: Workspace, agent: WorkspaceAgent): Long? {
+        val stack = ArrayDeque<ProcessHandle>(ProcessHandle.current().children().toList())
+        while (stack.isNotEmpty()) {
+            val processHandle = stack.removeLast()
+            val cmdLine = processHandle.info().commandLine().getOrNull()
+            ctx.logger.debug("SSH command PID: ${processHandle.pid()} Command: $cmdLine")
+            if (cmdLine != null && cmdLine.isSshCommandFor(ws, agent)) {
+                ctx.logger.debug("SSH command with PID: ${processHandle.pid()} and Command: $cmdLine matches ${ws.name}.${agent.name}")
+                return processHandle.pid()
+            } else {
+                stack.addAll(processHandle.children().toList())
+            }
+        }
+        return null
+    }
+
+    private fun String.isSshCommandFor(ws: Workspace, agent: WorkspaceAgent): Boolean {
+        // usage-app is present only in the ProxyCommand
+        return !this.contains("--usage-app=jetbrains") && this.contains("${ws.name}.${agent.name}")
+    }
+}

src/main/kotlin/com/coder/toolbox/sdk/v2/models/NetworkMetrics.kt

@@ -0,0 +1,41 @@
+package com.coder.toolbox.sdk.v2.models
+
+import com.squareup.moshi.Json
+import com.squareup.moshi.JsonClass
+
+/**
+ * Coder ssh network metrics. All properties are optional
+ * because Coder Connect only populates `using_coder_connect`
+ * while p2p doesn't populate this property.
+ */
+@JsonClass(generateAdapter = true)
+data class NetworkMetrics(
+    @Json(name = "p2p")
+    val p2p: Boolean?,
+
+    @Json(name = "latency")
+    val latency: Double?,
+
+    @Json(name = "preferred_derp")
+    val preferredDerp: String?,
+
+    @Json(name = "derp_latency")
+    val derpLatency: Map<String, Double>?,
+
+    @Json(name = "upload_bytes_sec")
+    val uploadBytesSec: Long?,
+
+    @Json(name = "download_bytes_sec")
+    val downloadBytesSec: Long?,
+
+    @Json(name = "using_coder_connect")
+    val usingCoderConnect: Boolean?
+) {
+    fun toPretty(): String {
+        return if (p2p == true) {
+            "Direct (${latency}ms) \u00B7 Download \u2193 $downloadBytesSec b/s \u00B7 Upload \u2191 $uploadBytesSec b/s"
+        } else {
+            "$preferredDerp (${latency}ms) \u00B7 Download \u2193 $downloadBytesSec b/s \u00B7 Upload \u2191 $uploadBytesSec b/s"
+        }
+    }
+}

src/main/kotlin/com/coder/toolbox/settings/ReadOnlyCoderSettings.kt

@@ -110,6 +110,12 @@ interface ReadOnlyCoderSettings {
      */
     val sshConfigOptions: String?
 
+
+    /**
+     * The path where network information for SSH hosts are stored
+     */
+    val networkInfoDir: String
+
     /**
      * The default URL to show in the connection window.
      */

src/main/kotlin/com/coder/toolbox/store/CoderSettingsStore.kt

@@ -65,6 +65,11 @@ class CoderSettingsStore(
     override val sshLogDirectory: String? get() = store[SSH_LOG_DIR]
     override val sshConfigOptions: String?
         get() = store[SSH_CONFIG_OPTIONS].takeUnless { it.isNullOrEmpty() } ?: env.get(CODER_SSH_CONFIG_OPTIONS)
+    override val networkInfoDir: String
+        get() = store[NETWORK_INFO_DIR].takeUnless { it.isNullOrEmpty() } ?: getDefaultGlobalDataDir()
+            .resolve("ssh-network-metrics")
+            .normalize()
+            .toString()
 
     /**
      * The default URL to show in the connection window.
@@ -232,6 +237,10 @@ class CoderSettingsStore(
         store[SSH_LOG_DIR] = path
     }
 
+    fun updateNetworkInfoDir(path: String) {
+        store[NETWORK_INFO_DIR] = path
+    }
+
     fun updateSshConfigOptions(options: String) {
         store[SSH_CONFIG_OPTIONS] = options
     }

src/main/kotlin/com/coder/toolbox/store/StoreKeys.kt

@@ -38,3 +38,5 @@ internal const val SSH_LOG_DIR = "sshLogDir"
 
 internal const val SSH_CONFIG_OPTIONS = "sshConfigOptions"
 
+internal const val NETWORK_INFO_DIR = "networkInfoDir"
+

src/main/kotlin/com/coder/toolbox/views/CoderSettingsPage.kt

@@ -56,6 +56,8 @@ class CoderSettingsPage(context: CoderToolboxContext, triggerSshConfig: Channel<
         TextField(context.i18n.ptrl("Extra SSH options"), settings.sshConfigOptions ?: "", TextType.General)
     private val sshLogDirField =
         TextField(context.i18n.ptrl("SSH proxy log directory"), settings.sshLogDirectory ?: "", TextType.General)
+    private val networkInfoDirField =
+        TextField(context.i18n.ptrl("SSH network metrics directory"), settings.networkInfoDir, TextType.General)
 
 
     override val fields: StateFlow<List<UiField>> = MutableStateFlow(
@@ -73,6 +75,7 @@ class CoderSettingsPage(context: CoderToolboxContext, triggerSshConfig: Channel<
             disableAutostartField,
             enableSshWildCardConfig,
             sshLogDirField,
+            networkInfoDirField,
             sshExtraArgs,
         )
     )
@@ -104,6 +107,7 @@ class CoderSettingsPage(context: CoderToolboxContext, triggerSshConfig: Channel<
                     }
                 }
                 context.settingsStore.updateSshLogDir(sshLogDirField.textState.value)
+                context.settingsStore.updateNetworkInfoDir(networkInfoDirField.textState.value)
                 context.settingsStore.updateSshConfigOptions(sshExtraArgs.textState.value)
             }
         )

src/main/resources/localization/defaultMessages.po

@@ -128,4 +128,10 @@ msgid "Extra SSH options"
 msgstr ""
 
 msgid "SSH proxy log directory"
+msgstr ""
+
+msgid "SSH network metrics directory"
+msgstr ""
+
+msgid "Network Metrics"
 msgstr ""

src/test/kotlin/com/coder/toolbox/cli/CoderCLIManagerTest.kt

@@ -18,6 +18,7 @@ import com.coder.toolbox.store.DISABLE_AUTOSTART
 import com.coder.toolbox.store.ENABLE_BINARY_DIR_FALLBACK
 import com.coder.toolbox.store.ENABLE_DOWNLOADS
 import com.coder.toolbox.store.HEADER_COMMAND
+import com.coder.toolbox.store.NETWORK_INFO_DIR
 import com.coder.toolbox.store.SSH_CONFIG_OPTIONS
 import com.coder.toolbox.store.SSH_CONFIG_PATH
 import com.coder.toolbox.store.SSH_LOG_DIR
@@ -510,7 +511,10 @@ internal class CoderCLIManagerTest {
                         HEADER_COMMAND to it.headerCommand,
                         SSH_CONFIG_PATH to tmpdir.resolve(it.input + "_to_" + it.output + ".conf").toString(),
                         SSH_CONFIG_OPTIONS to it.extraConfig,
-                        SSH_LOG_DIR to (it.sshLogDirectory?.toString() ?: "")
+                        SSH_LOG_DIR to (it.sshLogDirectory?.toString() ?: ""),
+                        NETWORK_INFO_DIR to tmpdir.parent.resolve("coder-toolbox")
+                            .resolve("ssh-network-metrics")
+                            .normalize().toString()
                     ),
                     env = it.env,
                     context.logger,
@@ -531,6 +535,7 @@ internal class CoderCLIManagerTest {
 
             // Output is the configuration we expect to have after configuring.
             val coderConfigPath = ccm.localBinaryPath.parent.resolve("config")
+            val networkMetricsPath = tmpdir.parent.resolve("coder-toolbox").resolve("ssh-network-metrics")
             val expectedConf =
                 Path.of("src/test/resources/fixtures/outputs/").resolve(it.output + ".conf").toFile().readText()
                     .replace(newlineRe, System.lineSeparator())
@@ -539,6 +544,10 @@ internal class CoderCLIManagerTest {
                         "/tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64",
                         escape(ccm.localBinaryPath.toString())
                     )
+                    .replace(
+                        "/tmp/coder-toolbox/ssh-network-metrics",
+                        escape(networkMetricsPath.toString())
+                    )
                     .let { conf ->
                         if (it.sshLogDirectory != null) {
                             conf.replace("/tmp/coder-toolbox/test.coder.invalid/logs", it.sshLogDirectory.toString())

src/test/resources/fixtures/outputs/append-blank-newlines.conf

@@ -4,7 +4,7 @@
 
 # --- START CODER JETBRAINS TOOLBOX test.coder.invalid
 Host coder-jetbrains-toolbox--owner--foo.agent1--test.coder.invalid
-  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --usage-app=jetbrains owner/foo.agent1
+  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --network-info-dir /tmp/coder-toolbox/ssh-network-metrics --usage-app=jetbrains owner/foo.agent1
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

src/test/resources/fixtures/outputs/append-blank.conf

@@ -1,6 +1,6 @@
 # --- START CODER JETBRAINS TOOLBOX test.coder.invalid
 Host coder-jetbrains-toolbox--owner--foo.agent1--test.coder.invalid
-  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --usage-app=jetbrains owner/foo.agent1
+  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --network-info-dir /tmp/coder-toolbox/ssh-network-metrics --usage-app=jetbrains owner/foo.agent1
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

src/test/resources/fixtures/outputs/append-no-blocks.conf

@@ -5,7 +5,7 @@ Host test2
 
 # --- START CODER JETBRAINS TOOLBOX test.coder.invalid
 Host coder-jetbrains-toolbox--owner--foo.agent1--test.coder.invalid
-  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --usage-app=jetbrains owner/foo.agent1
+  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --network-info-dir /tmp/coder-toolbox/ssh-network-metrics --usage-app=jetbrains owner/foo.agent1
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

src/test/resources/fixtures/outputs/append-no-newline.conf

@@ -4,7 +4,7 @@ Host test2
   Port 443
 # --- START CODER JETBRAINS TOOLBOX test.coder.invalid
 Host coder-jetbrains-toolbox--owner--foo.agent1--test.coder.invalid
-  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --usage-app=jetbrains owner/foo.agent1
+  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --network-info-dir /tmp/coder-toolbox/ssh-network-metrics --usage-app=jetbrains owner/foo.agent1
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

src/test/resources/fixtures/outputs/append-no-related-blocks.conf

@@ -11,7 +11,7 @@ some jetbrains config
 
 # --- START CODER JETBRAINS TOOLBOX test.coder.invalid
 Host coder-jetbrains-toolbox--owner--foo.agent1--test.coder.invalid
-  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --usage-app=jetbrains owner/foo.agent1
+  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --network-info-dir /tmp/coder-toolbox/ssh-network-metrics --usage-app=jetbrains owner/foo.agent1
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

src/test/resources/fixtures/outputs/disable-autostart.conf

@@ -1,6 +1,6 @@
 # --- START CODER JETBRAINS TOOLBOX test.coder.invalid
 Host coder-jetbrains-toolbox--owner--foo.agent1--test.coder.invalid
-  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --disable-autostart --usage-app=jetbrains owner/foo.agent1
+  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --disable-autostart --network-info-dir /tmp/coder-toolbox/ssh-network-metrics --usage-app=jetbrains owner/foo.agent1
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

src/test/resources/fixtures/outputs/extra-config.conf

@@ -1,6 +1,6 @@
 # --- START CODER JETBRAINS TOOLBOX test.coder.invalid
 Host coder-jetbrains-toolbox--owner--foo.agent1--test.coder.invalid
-  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --usage-app=jetbrains owner/foo.agent1
+  ProxyCommand /tmp/coder-toolbox/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-toolbox/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --network-info-dir /tmp/coder-toolbox/ssh-network-metrics --usage-app=jetbrains owner/foo.agent1
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null