;; Fibers: cooperative, event-driven user-space threads.

;;;; Copyright (C) 2023, 2024 Free Software Foundation, Inc.
;;;;
;;;; This library is free software; you can redistribute it and/or
;;;; modify it under the terms of the GNU Lesser General Public
;;;; License as published by the Free Software Foundation; either
;;;; version 3 of the License, or (at your option) any later version.
;;;;
;;;; This library is distributed in the hope that it will be useful,
;;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;;; Lesser General Public License for more details.
;;;;
;;;; You should have received a copy of the GNU Lesser General Public License
;;;; along with this program.  If not, see <http://www.gnu.org/licenses/>.
;;;;

;;;; Hierarchical timer wheel inspired by Juho Snellman's "Ratas".  For a
;;;; detailed discussion, see:
;;;;
;;;;   https://www.snellman.net/blog/archive/2016-07-27-ratas-hierarchical-timer-wheel/
;;;;
;;;; Ported from
;;;; https://github.com/snabbco/snabb/blob/master/src/lib/fibers/timer.lua,
;;;; by Andy Wingo.

(define-module (fibers timer-wheel)
  #:use-module (srfi srfi-9)
  #:use-module (ice-9 match)
  #:use-module (ice-9 format)
  #:export (make-timer-wheel
            timer-wheel-add!
            timer-wheel-remove!
            timer-wheel-next-entry-time
            timer-wheel-next-tick-start
            timer-wheel-next-tick-end
            timer-wheel-advance!
            timer-wheel-dump))

(define *slots* 256)
(define *slots-bits* 8)
(define *slots-mask* 255)

(define-record-type <timer-entry>
  (make-timer-entry prev next time obj)
  timer-entry?
  (prev timer-entry-prev set-timer-entry-prev!)
  (next timer-entry-next set-timer-entry-next!)
  (time timer-entry-time set-timer-entry-time!)
  (obj timer-entry-obj   set-timer-entry-obj!))

(define-record-type <timer-wheel>
  (%make-timer-wheel time-base shift cur slots outer next-entry-time)
  timer-wheel?
  (time-base timer-wheel-time-base set-timer-wheel-time-base!)
  (shift timer-wheel-shift)
  (cur timer-wheel-cur set-timer-wheel-cur!)
  (slots timer-wheel-slots)
  (outer timer-wheel-outer set-timer-wheel-outer!)
  (next-entry-time %timer-wheel-next-entry-time
                   set-timer-wheel-next-entry-time!))

(define (push-timer-entry! entry head)
  (match head
    (($ <timer-entry> prev)
     (set-timer-entry-prev! entry prev)
     (set-timer-entry-next! entry head)
     (set-timer-entry-prev! head entry)
     (set-timer-entry-next! prev entry))))

(define (make-slots)
  (let ((slots (make-vector *slots* #f)))
    (let lp ((i 0))
      (when (< i *slots*)
        (let ((entry (make-timer-entry #f #f #f #f)))
          (set-timer-entry-prev! entry entry)
          (set-timer-entry-next! entry entry)
          (vector-set! slots i entry))
        (lp (1+ i))))
    slots))

(define (time->slot-index internal-time shift)
  (logand (ash internal-time (- shift)) *slots-mask*))

(define (compute-time-base shift internal-time)
  (logand internal-time (lognot (1- (ash 1 shift)))))

(define* (make-timer-wheel #:key (now (get-internal-real-time))
                           ;; Default to at-least-millisecond precision.
                           (precision 1000)
                           (shift (let lp ((shift 0))
                                    (if (< (ash internal-time-units-per-second
                                                (- (1+ shift)))
                                           precision)
                                        shift
                                        (lp (1+ shift))))))
  (%make-timer-wheel (compute-time-base (+ shift *slots-bits*) now)
                     shift
                     (time->slot-index now shift)
                     (make-slots)
                     #f
                     #f))

(define (add-outer-wheel! inner)
  (match inner
    (($ <timer-wheel> time-base shift cur slots #f)
     (let* ((next-outer-tick (+ time-base (ash *slots* shift)))
            (outer (make-timer-wheel #:now next-outer-tick
                                     #:shift (+ shift *slots-bits*))))
       (set-timer-wheel-outer! inner outer)
       outer))))

(define (next-tick-start time-base cur shift)
  (+ time-base (ash cur shift)))
(define (next-tick-end time-base cur shift)
  (next-tick-start time-base (1+ cur) shift))

(define (timer-wheel-next-tick-start wheel)
  (match wheel
    (($ <timer-wheel> time-base shift cur slots outer)
     (next-tick-start time-base cur shift))))
(define (timer-wheel-next-tick-end wheel)
  (match wheel
    (($ <timer-wheel> time-base shift cur slots outer)
     (next-tick-end time-base cur shift))))

(define (timer-wheel-add! wheel t obj)
  (match wheel
    (($ <timer-wheel> time-base shift cur slots outer next-entry-time)
     (unless (eq? next-entry-time 'unknown)
       (when (or (not next-entry-time) (< t next-entry-time))
         (set-timer-wheel-next-entry-time! wheel t)))
     (let ((offset (ash (- t (next-tick-start time-base cur shift))
                        (- shift))))
       (cond
        ((< offset *slots*)
         (let ((idx (logand (+ cur (max offset 0)) *slots-mask*))
               (entry (make-timer-entry #f #f t obj)))
           (push-timer-entry! entry (vector-ref slots idx))
           entry))
        (else
         (timer-wheel-add! (or outer (add-outer-wheel! wheel)) t obj)))))))

(define (timer-wheel-remove! wheel entry)
  "Remove @var{entry}, a timer entry as returned by @code{timer-wheel-add!},
from @var{wheel}."
  (define (invalidate-next-entry-time! wheel)
    (set-timer-wheel-next-entry-time! wheel 'unknown)
    (let ((outer (timer-wheel-outer wheel)))
      (when outer
        (invalidate-next-entry-time! outer))))

  (match entry
    (($ <timer-entry> prev next)
     (set-timer-entry-next! prev next)
     (set-timer-entry-prev! entry entry)

     (set-timer-entry-prev! next prev)
     (set-timer-entry-next! entry entry)

     (set-timer-entry-obj! entry #f)
     (set-timer-entry-time! entry #f)

     ;; Removing ENTRY might have change the next entry time for WHEEL.
     ;; Thus invalidate it so that it gets recomputed.
     (invalidate-next-entry-time! wheel))))

(define (timer-wheel-next-entry-time wheel)
  (define (slot-min-time head)
    (let lp ((entry (timer-entry-next head)) (min #f))
      (if (eq? entry head)
          min
          (match entry
            (($ <timer-entry> prev next t obj)
             (lp next (if (and min (< min t)) min t)))))))
  (define (compute-next-entry-time cur slots outer)
    (let lp ((i 0))
      (cond
       ((< i *slots*)
        (match (slot-min-time
                (vector-ref slots (logand (+ cur i) *slots-mask*)))
          (#f (lp (1+ i))) ;; Empty slot.
          (t
           ;; Unless we just migrated entries from outer to inner wheel
           ;; on the last tick, outer wheel overlaps with inner.
           (let ((outer-t (match outer
                            (#f #f)
                            (($ <timer-wheel> time-base shift cur slots outer)
                             (slot-min-time (vector-ref slots cur))))))
             (if outer-t
                 (min t outer-t)
                 t)))))
       (outer (timer-wheel-next-entry-time outer))
       (else #f))))
  (match wheel
    (($ <timer-wheel> time-base shift cur slots outer next-entry-time)
     (match next-entry-time
       ('unknown
        (let ((t (compute-next-entry-time cur slots outer)))
          (set-timer-wheel-next-entry-time! wheel t)
          t))
       (t t)))))

(define* (timer-wheel-dump wheel #:key (port (current-output-port))
                           (level 0)
                           (process-time
                            (lambda (t)
                              (/ t 1.0 internal-time-units-per-second))))
  (match wheel
    (($ <timer-wheel> time-base shift cur slots outer)
     (let ((start (next-tick-start time-base cur shift)))
       (let lp ((i 0))
         (when (< i *slots*)
           (let* ((head (vector-ref slots (logand *slots-mask* (+ cur i))))
                  (entry (timer-entry-next head)))
             (unless (eq? entry head)
               (format port "level ~a, tick +~a (~a-~a):\n" level i
                       (process-time (+ start (ash i shift)))
                       (process-time (+ start (ash (1+ i) shift))))
               (let lp ((entry entry))
                 (match entry
                   (($ <timer-entry> _ next t obj)
                    (format port "  ~a: ~a\n" (process-time t) obj)
                    (unless (eq? next head)
                      (lp next)))))))
           (lp (1+ i)))))
     (when outer
       (timer-wheel-dump outer #:port port #:level (1+ level)
                         #:process-time process-time)))))

(define (timer-wheel-advance! wheel t schedule!)
  (define (tick!)
    ;; Define as syntax to make sure it gets inlined; otherwise the
    ;; compiler currently ends up making a closure.
    (define-syntax-rule (advance-wheel! wheel visit-timer-entry!)
      (match wheel
        (($ <timer-wheel> time-base shift cur slots outer)
         (let ((head (vector-ref slots cur)))
           (let lp ()
             (match head
               (($ <timer-entry> _ entry)
                (cond
                 ((eq? entry head) #f)
                 (else
                  (match entry
                    (($ <timer-entry> _ next t obj)
                     (set-timer-wheel-next-entry-time! wheel 'unknown)
                     (set-timer-entry-next! head next)
                     (set-timer-entry-prev! next head)
                     (visit-timer-entry! entry t obj)
                     (lp)))))))))
         (let ((cur (logand (1+ cur) *slots-mask*)))
           (set-timer-wheel-cur! wheel cur)
           (when (zero? cur)
             (let ((time-base (+ time-base (ash *slots* shift))))
               (set-timer-wheel-time-base! wheel time-base)
               (when outer (tick-outer! wheel outer))))))))

    (define (tick-outer! inner outer)
      (define (add-to-inner! entry t obj)
        (match inner
          (($ <timer-wheel> time-base shift cur slots outer)
           (let ((new-head (vector-ref slots (time->slot-index t shift))))
             (push-timer-entry! entry new-head)))))
      (advance-wheel! outer add-to-inner!))

    (advance-wheel! wheel (lambda (entry t obj) (schedule! obj))))

  (match wheel
    (($ <timer-wheel> time-base shift cur slots outer)
     (let ((inc (ash 1 shift)))
       (let lp ((next-tick-end (next-tick-end time-base cur shift)))
         (when (<= next-tick-end t)
           (tick!)
           (lp (+ next-tick-end inc))))))))

