#lang scheme/base

;; The core profiler sample collector
;; (This module is a private tool for collecting profiling data, and should not
;; be used as is.)

(provide create-sampler)

;; create-sampler : creates a sample collector thread, which tracks the given
;; `to-track' value every `delay' seconds.
;; * The input value can be either a thread (track just that thread), a
;;   custodian (track all threads managed by the custodian), or a list of
;;   threads and/or custodians.  If a custodian is given, it must be
;;   subordinate to the current custodian or to the given super custodian.
;; * The collected values are (<thread-id> <thread-time> . <stack>), where
;;   - The <thread-id> is an integer number identifying the thread, starting
;;     from 0.  If the collected data has thread ids in a 0..N range
;;     (exclusive) , then there were N threads that were observed.  (This can
;;     be relevant when tracking a custodian, where threads can change
;;     dynamically.)
;;   - The <thread-time> is the result of `current-process-milliseconds' for
;;     the thread that this sample came from.  Note that these numbers will not
;;     start at 0, since threads are likely to run before a sample is
;;     collected.
;;   - Finally, the <stack> part is a snapshot of the thread's stack, as
;;     grabbed by `continuation-mark-set->context'.  The values in these
;;     snapshots are interned to reduce memory load.
;;   The results are collected sequentially, so they're always sorted from the
;;   newest to the oldest.  Remember that these results should be considered
;;   private for the profiler collection, and can change when more information
;;   needs to be collected.
;; * Returns a "controller" function that accepts messages to control the
;;   sampler thread.  The current set of messages that the controller
;;   understands are:
;;   - 'pause and 'resume: stops or resumes collecting samples.  These messages
;;     can be nested.  Note that the thread will continue running it just won't
;;     collect snapshots.
;;   - 'stop: kills the collector thread.  Should be called when you have your
;;     data.  (There is no message to start a new sampler thread, although
;;     adding one will not be difficult.)
;;   - 'set-tracked! <new>: changes the thread/s and/or custodian/s to track.
;;     (Custodians should still be subordinate to the original one or to the
;;     given argument.)
;;   - 'set-delay! <new>: changes the sampling delay.  This means that we won't
;;     have a direct correlation between the number of samples and the time
;;     they represent -- but the samples are statistical snapshots anyway, and
;;     the results are not formulated in terms of time spent.  (The time spent
;;     could be added of course, but it is best to do that in terms of the
;;     start/stop times)
;;   - 'get-snapshots: returns the currently collected list of snapshots.  Note
;;     that this can be called multiple times, each will return the data that
;;     is collected up to that point in time.
(define (create-sampler to-track delay [super-cust (current-custodian)])
  ;; the collected data
  (define snapshots '())
  ;; intern the entries (which are (cons id/#f srcloc/#f))
  (define entry-table (make-hash))
  (define (intern-entry entry)
    (let* ([key (or (cdr entry) (car entry))]
           [en  (hash-ref entry-table key #f)])
      (if en
        ;; ELI: is this sanity check needed?
        ;; (if (equal? en entry)
        ;;   en
        ;;   (error 'profile "internal error: assumption invalid"))
        en
        (begin (hash-set! entry-table key entry) entry))))
  (define (validate to-track who)
    (let loop ([t to-track])
      (cond
        [(thread? t)]
        [(list? t) (for-each loop t)]
        [(not (custodian? t))
         (raise-type-error
          who "thread, custodian, or a list of threads/csutodians" to-track)]
        ;; test that it's subordinate
        [(with-handlers ([exn:fail:contract? (lambda (_) #t)])
           (custodian-managed-list t super-cust) #f)
         (error who "got an insubordinate custodian")])))
  (define paused 0)
  (define thread-id
    (let ([next-id 0] [t (make-weak-hasheq)])
      (lambda (thd)
        (or (hash-ref t thd #f)
            (let ([id next-id])
              (set! next-id (add1 next-id))
              (hash-set! t thd id)
              id)))))
  (define (sampler)
    (sleep delay)
    (when (paused . <= . 0)
      (let loop ([t to-track])
        (cond [(thread? t)
               (unless (eq? t sampler-thread)
                 (set! snapshots
                       (cons (list* (thread-id t)
                                    (current-process-milliseconds t)
                                    (map intern-entry
                                         (continuation-mark-set->context
                                          (continuation-marks t))))
                             snapshots)))]
              [(custodian? t)
               (for-each loop (custodian-managed-list t super-cust))]
              ;; cannot assume that it's a list: we might get other values from
              ;; a custodian managed list
              [(list? t) (for-each loop t)])))
    (sampler))
  (define cpu-time   0)
  (define start-time (current-process-milliseconds))
  (define (add-time)
    (when (paused . <= . 0)
      (let ([cur (current-process-milliseconds)])
        (set! cpu-time (+ cpu-time (- cur start-time)))
        (set! start-time cur))))
  (define (ignore-time)
    (when (paused . <= . 0)
      (set! start-time (current-process-milliseconds))))
  (define sampler-thread
    (begin (validate to-track 'create-sampler)
           (thread sampler)))
  ;; use a sema to avoid mutations from different threads, the sampler thread
  ;; is only reading these values so it doesn't use it.
  (define sema (make-semaphore 1))
  (define (sampler-controller msg [arg #f])
    (define-syntax-rule (w/sema body ...)
      (call-with-semaphore sema (lambda () body ...)))
    (case msg
      [(pause)  (w/sema (add-time) (set! paused (add1 paused)))]
      [(resume) (w/sema (set! paused (sub1 paused)) (ignore-time))]
      [(stop)   (kill-thread sampler-thread) (add-time) (set! paused +inf.0)]
      [(set-tracked!) (validate arg 'sampler-controller)
                      (w/sema (set! to-track arg))]
      [(set-delay!)   (w/sema (set! delay arg))]
      [(get-snapshots) (add-time) (cons cpu-time snapshots)]
      [else (error 'sampler-controller "unknown message: ~e" msg)]))
  sampler-controller)