foo.rs

#![feature(phase)]

#[phase(syntax, link)]
extern crate green;
extern crate rtinstrument;

green_start!(main)

fn main() {
    let msgs = rtinstrument::instrument::<green::task::GreenTask>(|| {
        let (tx, rx) = channel();
        for _ in range(0, 10) {
            let tx = tx.clone();
            spawn(proc() {
                println!("baz");
                tx.send(());
            });
        }
        for _ in range(0, 10) {
            rx.recv();
        }
    });

    for msg in msgs.iter() {
        println!("{}", msg);
    }
}

Makefile

all:
	rustc rtinstrument.rs
	rustc foo.rs -L .
	./foo | racket rust-visualizer.rkt &

rtinstrument.rs

#![crate_type = "rlib"]

extern crate time;
extern crate sync;

use sync::Arc;
use std::any::Any;
use std::mem;
use std::raw;
use std::rt::Runtime;
use std::rt::exclusive::Exclusive;
use std::rt::local::Local;
use std::rt::rtio;
use std::rt::task::{Task, TaskOpts, BlockedTask};

#[deriving(Show, Clone)]
pub struct Message {
    /// Time at which this event happened
    pub timestamp: u64,
    /// Rust task that performed this event
    pub task_id: uint,
    /// OS thread that performed this event
    pub thread_id: uint,
    /// Short string description of the event that happened
    pub desc: String,
}

struct InstrumentedRuntime<R> {
    inner: Option<Box<R>>,
    messages: Arc<Exclusive<Vec<Message>>>,
    thread_id: uint,
}

/// Instrument all code run inside the specific block, returning a vector of all
/// messages which occurred.
pub fn instrument<R: 'static + Runtime + Send>(f: ||) -> Vec<Message> {
    install::<R>(Arc::new(Exclusive::new(Vec::new())));
    f();
    let rt = uninstall::<R>();
    unsafe { rt.messages.lock().clone() }
}

/// Installs an instrumented runtime which will append to the given vector of
/// messages.
///
/// The instrumented runtime is installed into the current task.
fn install<R: 'static + Runtime + Send>(messages: Arc<Exclusive<Vec<Message>>>) {
    let mut task = Local::borrow(None::<Task>);
    let mut rt = task.maybe_take_runtime::<R>().unwrap();
    let new_rt = box InstrumentedRuntime {
        thread_id: rt.local_io().map(|mut i| {
            let i: raw::TraitObject = unsafe { mem::transmute(i.get()) };
            i.data as uint
        }).unwrap_or(0),
        inner: Some(rt),
        messages: messages
    };
    new_rt.log("spawn");
    task.put_runtime(new_rt);
}

/// Uninstalls the runtime from the current task, returning the instrumented
/// runtime.
fn uninstall<R: 'static + Runtime + Send>() -> InstrumentedRuntime<R> {
    let mut task = Local::borrow(None::<Task>);
    let mut rt = task.maybe_take_runtime::<InstrumentedRuntime<R>>().unwrap();
    task.put_runtime(rt.inner.take().unwrap());
    rt.log("death");
    *rt
}

impl<R: 'static + Runtime + Send>  InstrumentedRuntime<R> {
    /// Puts this runtime back into the local task
    fn put(mut ~self) {
        assert!(self.inner.is_none());

        let mut task: Box<Task> = Local::take();
        let rt = task.maybe_take_runtime::<R>().unwrap();
        self.inner = Some(rt);
        task.put_runtime(self);
        Local::put(task);
    }

    /// Logs a message into this runtime
    fn log(&self, msg: &str) {
        let mut messages = unsafe { self.messages.lock() };
        messages.push(Message {
            timestamp: time::precise_time_ns(),
            desc: msg.to_str(),
            task_id: self as *_ as uint,
            thread_id: self.thread_id,
        });
    }
}

impl<R: 'static + Runtime + Send> Runtime for InstrumentedRuntime<R> {
    fn yield_now(mut ~self, cur_task: Box<Task>) {
        self.log("yield");
        self.inner.take().unwrap().yield_now(cur_task);
        self.put()
    }

    fn maybe_yield(mut ~self, cur_task: Box<Task>) {
        self.log("maybe-yield");
        self.inner.take().unwrap().maybe_yield(cur_task);
        self.put()
    }

    fn deschedule(mut ~self, times: uint, cur_task: Box<Task>,
                  f: |BlockedTask| -> Result<(), BlockedTask>) {
        self.log("deschedule");
        self.inner.take().unwrap().deschedule(times, cur_task, f);
        self.log("wakeup");
        self.put()
    }

    fn reawaken(mut ~self, to_wake: Box<Task>) {
        // TODO: log which task awoke which task?
        self.inner.take().unwrap().yield_now(to_wake);
        self.put()
    }

    fn spawn_sibling(mut ~self,
                     cur_task: Box<Task>,
                     opts: TaskOpts,
                     f: proc():Send) {
        // Be sure to install an instrumented runtime for the spawned sibling by
        // specifying a new runtime.
        let messages = self.messages.clone();
        self.inner.take().unwrap().spawn_sibling(cur_task, opts, proc() {
            install::<R>(messages);
            f();
            drop(uninstall::<R>());
        });
        self.put()
    }

    fn local_io<'a>(&'a mut self) -> Option<rtio::LocalIo<'a>> {
        self.inner.get_mut_ref().local_io()
    }

    fn stack_bounds(&self) -> (uint, uint) {
        self.inner.get_ref().stack_bounds()
    }

    fn can_block(&self) -> bool {
        self.inner.get_ref().can_block()
    }

    // FIXME: This is a serious code smell and this should not exist at all.
    fn wrap(~self) -> Box<Any> { self as Box<Any> }
}

rust-visualizer.rkt

#lang racket

(require future-visualizer future-visualizer/trace)

(define log-regexp
  "^Message { timestamp: ([0-9]+), task_id: ([0-9]+), thread_id: ([0-9]+), desc: (.+) }")

(define index 0)
(define (next-index!) (begin0 index (set! index (add1 index))))

;; parse-event : string? -> (listof indexed-future-event?)
;; maybe generate 0, 1 or more future events from one rust event
(define (parse-event s)
  (match (regexp-match log-regexp s)
    [(list _ timestamp-string task-id-string thread-id-string desc)
     (define timestamp (/ (string->number timestamp-string)
                          ;; to be on the same scale as future timestamps
                          1000000.0))
     (define task-id   (string->number task-id-string))
     (define thread-id (string->number thread-id-string))
     (and timestamp task-id thread-id
          (build-indexed-events timestamp task-id thread-id desc))]
    ;; probably some random program output, ignore
    [_ '()]))

;; build-indexed-events : integer? integer? integer? string?
;;                         -> (listof indexed-future-event?)
;; the index serves to disambiguate order in case of identical timestamps
;; (struct future-event (future-id proc-id action time prim-name user-data))
(define (build-indexed-events timestamp task-id thread-id desc)
  (define kinds (desc->kinds desc))
  (define (make kind)
    (indexed-future-event
     (next-index!)
     (future-event (match kind
                     ['create #f] ; seems to be part of the protocol
                     [_ task-id])
                   thread-id
                   kind
                   timestamp
                   #f
                   (match kind ; user-data field
                     ['create task-id]
                     ;; none of the other messages we produce need user data
                     [_ #f]))))
  (map make kinds))

;; desc->kinds : string? -> (listof symbol?)
;; returns a list of the kinds of future events that should be emitted
;; based on the rust event observed
(define (desc->kinds desc)
  (match desc
    ["spawn"       '(create start-work)]
    ["death"       '(complete)]
    ["yield"       '(sync)]
    ["maybe-yield" '(sync)]
    ["deschedule"  '(sync)]
    ["wakeup"      '(start-work)]))



(module+ main
  (define events (append-map parse-event (port->lines)))
  ;; (for-each displayln events)
  (show-visualizer #:timeline events)
  )