functiontrace-server 0.8.0

use functiontrace_server::function_trace::{FunctionTrace, TraceInitialization};
use functiontrace_server::profile_generator::{
    FirefoxProfile, FirefoxProfileThreadId, FirefoxThread,
};
use functiontrace_server::trace_streamer::TraceSender;
use serde::de::Deserialize;
use std::io::{BufReader, Read};
use std::path::Path;

use color_eyre::eyre::{eyre, Result, WrapErr};
use owo_colors::OwoColorize;
use tokio::io::AsyncReadExt;

#[cfg(feature = "debug-tracelog")]
use std::fs::File;
#[cfg(feature = "debug-tracelog")]
use std::io::prelude::*;
#[cfg(feature = "debug-tracelog")]
use std::time::Duration;

#[derive(argh::FromArgs)]
/** functiontrace-server <https://functiontrace.com>
The profile generation server for functiontrace.  This should rarely be run manually.
*/
struct Args {
    /// the directory to write generated profiles to
    #[argh(option, short = 'd')]
    directory: String,

    /// read debug output from a thread's previous run from the given file
    #[cfg(feature = "debug-tracelog")]
    #[argh(option)]
    debug_tracelog: Option<String>,
}

#[tokio::main]
async fn main() -> Result<()> {
    color_eyre::install()?;
    env_logger::Builder::from_default_env()
        .format_timestamp(None)
        .format_module_path(false)
        .init();
    let args: Args = argh::from_env();

    // Find the location we're supposed to output our profiles to.
    let output_dir = Path::new(&args.directory)
        .canonicalize()
        .wrap_err("The provided output directory doesn't exist")?;

    #[cfg(feature = "debug-tracelog")]
    if let Some(filename) = args.debug_tracelog {
        // We've been asked to read from a debug tracelog, rather than live input.
        log::info!("Parsing a previous set of inputs to functiontrace-server");

        // Attempt to parse the log, allowing us to make server-side changes to see the
        // impact on previous inputs.  We don't care about generating any real outputs
        // in this mode.
        let mut profile = FirefoxProfile::new(TraceInitialization {
            lang_version: "replayed".to_string(),
            platform: "replayed".to_string(),
            program_name: "replayed".to_string(),
            program_version: "replayed".to_string(),
            time: Duration::new(0, 0),
        });

        let tid = profile.register_thread();
        parse_thread_logs(File::open(filename)?, &tid)?;

        return Ok(());
    }

    // In the background, check if there's a more recent version of functiontrace available.
    let update_check = tokio::task::spawn(async {
        let latest = tokio::time::timeout(
            std::time::Duration::from_millis(500),
            check_latest::check_max_async!(),
        )
        .await;

        // We intentionally ignore all errors here since this isn't key functionality.
        latest.ok().map(|x| x.ok()).flatten().unwrap_or(None)
    });

    log::info!("Starting functiontrace-server");

    // Create a Unix socket to listen on for trace messages.
    let pipe = Path::new(&format!(
        "/tmp/functiontrace-server.sock.{}",
        std::process::id()
    ))
    .to_path_buf();

    if pipe.exists() {
        // This socket already exists - remove it so we can rebind.
        std::fs::remove_file(&pipe)?;
        log::warn!("Deleted an existing pipe - you should kill any other running instances of this application.");
    }

    let listener = tokio::net::UnixListener::bind(pipe).wrap_err("Race condition creating pipe")?;

    // The first message we're expecting is an encoded TraceInitialization.
    let mut profile = {
        let (mut socket, _addr) = listener.accept().await?;

        let info: TraceInitialization = {
            let mut buf = Vec::with_capacity(std::mem::size_of::<TraceInitialization>());
            socket
                .read_to_end(&mut buf)
                .await
                .wrap_err("Invalid initialization message")?;

            // TODO: We should print out what it failed on
            rmp_serde::from_slice(&buf).wrap_err("Failed to parse message")?
        };
        log::info!("Received a new trace connection: {:?}", info);

        FirefoxProfile::new(info)
    };

    let mut profile_threads = tokio::task::JoinSet::new();
    let mut client_tasks = tokio::task::JoinSet::new();
    let (mut spawned_tasks, mut retired_tasks) = (0, 0);

    loop {
        tokio::select! {
            // Listen to more connections, each of which corresponds to a client thread.
            connection = listener.accept() => {
                let mut client = connection.wrap_err("Error accepting a new client")?.0;

                // There's a new client ready to talk to us.  Register it, then spawn a task to
                // collect it's data and ferry it to blocking thread to be processed.
                let thread_info = profile.register_thread();
                log::info!("|{}| Connection from a new client", thread_info.tid);

                let mut sender = TraceSender::new(thread_info.tid);
                let receiver = sender.receiver();

                // Register this thread with the profile, then push the full log
                // parsing out into its own thread.
                profile_threads.spawn_blocking(move || {
                    parse_thread_logs(receiver, &thread_info)
                });

                spawned_tasks += 1;

                // Spawn a task to process the data this client sends us.
                client_tasks.spawn(async move {
                    let client_id = sender.client_id;

                    loop {
                        // Allocate a reasonably sized buffer to read data into.
                        let mut buf = Vec::new();
                        buf.resize(64 * 1024, 0);

                        match client.read(&mut buf).await {
                            Ok(0) => {
                                // Once we receive a wakeup for 0 bytes, we know that the other
                                // side disconnected.
                                let sent = sender.retire();
                                log::info!(
                                    "|{}| Retired client after reading {}",
                                    client_id,
                                    bytesize::to_string(sent as u64, false)
                                );
                                break;
                            }
                            Ok(len) => {
                                // We read some data, but there may be more.  Send this batch to
                                // the other consumer for now.
                                buf.truncate(len);
                                sender.send(buf);
                            }
                            Err(err) => {
                                if err.kind() == std::io::ErrorKind::WouldBlock {
                                    log::trace!("|{}| Blocking client", client_id);
                                    continue;
                                }

                                log::error!("|{}| Client disappeared: {}", client_id, err);
                                break;
                            }
                        }
                    }
                });
            },

            // One of our clients has disconnected and should be retired.
            Some(_) = client_tasks.join_next() => {
                retired_tasks += 1;

                if spawned_tasks == retired_tasks {
                    // We've registered clients and retired the same number that have registered.
                    // This means there are no clients left outstanding (and no others could be
                    // capable of talking to this socket), so we're done handling inbound messages.
                    break;
                }
            }
        };
    }

    // Wait for each of the threads to parse their traces, then add them to the profile.
    while let Some(t) = profile_threads.join_next().await {
        match t.expect("Parsing thread failed to join")? {
            ThreadTrace::ParsedThread(trace) => profile.finalize_thread(trace),
            // TODO: We should denote that this thread is corrupted in some fashion.
            // However, we've never observed this case in the wild, so don't know how it'll
            // manifest.
            ThreadTrace::CorruptedThread(trace) => profile.finalize_thread(trace),
        }
    }
    profile.export(output_dir)?;

    // Now that we've finished, print a notification if there's a newer version of functiontrace
    // for download.
    // NOTE: We don't want to print during execution since it might interleave poorly with other
    // text.
    if let Ok(Some(version)) = update_check.await {
        println!("{}", format!("functiontrace-server v{} has been released!  Consider updating to the newest release: {}", version, "cargo install functiontrace-server".italic()).bright_purple());
    }

    Ok(())
}

/// A valid thread's trace can either be parsed entirely successful or have corrupted data
/// somewhere in it.
pub enum ThreadTrace {
    /// The parsed thread logs.
    ParsedThread(FirefoxThread),
    /// The given thread failed to be parsed properly, likely due to a crash, and the trace is
    /// terminated unexpectedly.
    CorruptedThread(FirefoxThread),
}

/// Parses the inbound trace stream for the given thread, returning the parsed data.
fn parse_thread_logs(logs: impl Read, tid: &FirefoxProfileThreadId) -> Result<ThreadTrace> {
    let mut decoder = rmp_serde::decode::Deserializer::new(BufReader::new(logs));
    let id = tid.tid;

    #[cfg(feature = "debug-tracelog")]
    let debug_file = format!("functiontrace_raw.{}.dat", id);
    #[cfg(feature = "debug-tracelog")]
    let output_error = format!("Dumping output to {}.  This can be analyzed via the --debug-tracelog flag to functiontrace-server.", debug_file);

    // We should first receive a thread registration event.
    let mut thread = match FunctionTrace::deserialize(&mut decoder) {
        Ok(FunctionTrace::RegisterThread(registration)) => {
            log::info!("|{}| Parsing a new thread", id);
            FirefoxThread::new(registration, tid)
        }
        Ok(_) => {
            return Err(eyre!("Missing ThreadRegistration event for thread {}", id));
        }
        Err(e) => {
            #[cfg(feature = "debug-tracelog")]
            {
                // TODO: This is broken right now, as we aren't able to look backwards into the
                // decoder like this.  We should probably just buffer all output in this mode to a
                // Vec, allowing us to dump it if needed.
                log::error!("|{}| {}", id, output_error);
                File::create(debug_file)?.write_all(decoder.get_ref().get_ref())?;
            }

            return Err(e)
                .wrap_err_with(|| format!("Deserialization error initializing thread {}", id));
        }
    };

    // After the thread is registered, we receive various trace events until the socket is closed.
    loop {
        match FunctionTrace::deserialize(&mut decoder) {
            Ok(FunctionTrace::RegisterThread(registration)) => {
                log::error!(
                    "|{}| Received an unexpected registration event: {:?}",
                    id,
                    registration
                );
                break Err(eyre!("Unexpected ThreadRegistration"));
            }
            Ok(trace_log) => {
                log::trace!("|{}| {:?}", id, trace_log);
                thread.add_trace(trace_log);
            }
            Err(rmp_serde::decode::Error::InvalidMarkerRead(_)) => {
                // We get an invalid read once we're at the end of the loop.
                log::info!("|{}| Fully parsed thread!", id);
                break Ok(ThreadTrace::ParsedThread(thread));
            }
            Err(e) => {
                // We received an error, but we've parsed previous data that we don't want to drop.
                // Mark that this trace buffer was corrupted, and return the current thread state.
                //
                // XXX: We can't print the position of the deserializer because rmp only supports
                // `.position()` on std::io::Cursors for some reason.
                log::warn!("|{}| Deserialization error `{}` at (unknown offset)", id, e);

                #[cfg(feature = "debug-tracelog")]
                {
                    log::error!("|{}| {}", id, output_error);
                    File::create(debug_file)?.write_all(decoder.get_ref().get_ref())?;
                }

                break Ok(ThreadTrace::CorruptedThread(thread));
            }
        }
    }
}