rust: add option to kill server at stdin EOF (#4408)

Summary:
If `tensorboard` is to launch RustBoard as a subprocess rather than
relying on the user to launch it concurrently, then we should try hard
not to leave around a zombie server. An easy solution is an `atexit`
handler in Python TensorBoard, but this doesn’t handle the case when
TensorBoard crashes or gets SIGTERMed. On Linux, we can be notified when
our parent dies by setting the parent-death signal (`man 2 prctl`), but
this isn’t portable. On portable Unix, the child can poll its PPID to
see when it changes to `1` or the PID of some subreaper, but this isn’t
portable to Windows, which doesn’t update PPID when the parent dies.

Windows is not my strong suit, but some web searching didn’t turn up an
easy and clean solution portable to both Windows and Unix (any Windows
internals like “job objects” are a non-starter, since I can’t test
them). The one thing that I would expect to work everywhere is “just
wait until stdin closes and then exit”. If even that doesn’t work on
Windows, well, we can burn that bridge when we come to it.

Test Plan:
Build `bazel build //tensorboard/data/server`, then write a simple
Python driver:

```python
import subprocess
import time

server = "bazel-bin/tensorboard/data/server/server"
p = subprocess.Popen(
    [server, "--logdir", "/tmp/nonexistent", "-v", "--die-after-stdin"],
    stdin=subprocess.PIPE,
)
print(p.pid)
time.sleep(2)
```

Run it with `python test.py`, and note that after 2 seconds the server
prints “Stdin closed; exiting”. Run it again, and suspend (`^Z`) the
process before it finishes sleeping. Run `kill -SIGCONT CHILD_PID` with
the PID printed by the Python script to resume server execution; this is
just needed because your `^Z` propagates to all processes in the group.
Note that the server continues to print logs as it starts and finishes
new load cycles. Then, run `fg` and wait for the sleep to complete, and
note that the server again exits, as desired.

wchargin-branch: rust-die-after-stdin

Author: wchargin

tensorboard/data/server/main.rs

@@ -15,9 +15,11 @@ limitations under the License.
 
 use clap::Clap;
 use log::{debug, info, LevelFilter};
+use std::io::Read;
 use std::net::{IpAddr, SocketAddr};
 use std::path::PathBuf;
 use std::str::FromStr;
+use std::thread;
 use std::time::{Duration, Instant};
 use tokio::net::TcpListener;
 use tonic::transport::Server;
@@ -62,6 +64,16 @@ struct Opts {
     /// Use verbose output (-vv for very verbose output)
     #[clap(long = "verbose", short, parse(from_occurrences))]
     verbosity: u32,
+
+    /// Kill this server once stdin is closed
+    ///
+    /// While this server is running, read stdin to end of file and then kill the server. Used to
+    /// portably ensure that the server exits when the parent process dies, even due to a crash.
+    /// Don't set this if stdin is connected to a tty and the process will be backgrounded, since
+    /// then the server will receive `SIGTTIN` and its process will be stopped (in the `SIGSTOP`
+    /// sense) but not killed.
+    #[clap(long)]
+    die_after_stdin: bool,
 }
 
 /// A duration in seconds.
@@ -89,6 +101,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     });
     debug!("Parsed options: {:?}", opts);
 
+    if opts.die_after_stdin {
+        thread::Builder::new()
+            .name("StdinWatcher".to_string())
+            .spawn(die_after_stdin)
+            .expect("failed to spawn stdin watcher thread");
+    }
+
     let addr = SocketAddr::new(opts.host, opts.port);
     let listener = TcpListener::bind(addr).await?;
     let bound = listener.local_addr()?;
@@ -98,17 +117,20 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // leaks the outer commit structure (of constant size), not the pointers to the actual data.
     let commit: &'static Commit = Box::leak(Box::new(Commit::new()));
 
-    std::thread::spawn(move || {
-        let mut loader = LogdirLoader::new(commit, opts.logdir);
-        loop {
-            info!("Starting load cycle");
-            let start = Instant::now();
-            loader.reload();
-            let end = Instant::now();
-            info!("Finished load cycle ({:?})", end - start);
-            std::thread::sleep(opts.reload_interval.duration());
-        }
-    });
+    thread::Builder::new()
+        .name("Reloader".to_string())
+        .spawn(move || {
+            let mut loader = LogdirLoader::new(commit, opts.logdir);
+            loop {
+                info!("Starting load cycle");
+                let start = Instant::now();
+                loader.reload();
+                let end = Instant::now();
+                info!("Finished load cycle ({:?})", end - start);
+                thread::sleep(opts.reload_interval.duration());
+            }
+        })
+        .expect("failed to spawn reloader thread");
 
     let handler = DataProviderHandler { commit };
     Server::builder()
@@ -125,3 +147,12 @@ fn init_logging(default_log_level: LevelFilter) {
     use env_logger::{Builder, Env};
     Builder::from_env(Env::default().default_filter_or(default_log_level.to_string())).init();
 }
+
+/// Locks stdin and reads it to EOF, then exits the process.
+fn die_after_stdin() {
+    let stdin = std::io::stdin();
+    let stdin_lock = stdin.lock();
+    for _ in stdin_lock.bytes() {}
+    info!("Stdin closed; exiting");
+    std::process::exit(0);
+}