idk honestly
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

228 lines
8.8 KiB

  1. #lang racket
  2. ;; taskq:
  3. ;; infrastructure for a persistent task queue that guarantees each entry will execute to completion
  4. ;; at least once. entries are responsible for implementing idempotency in the case of interrupted
  5. ;; operations (due to a crash, for example)
  6. (require db crypto racket/fasl)
  7. (provide make-taskq taskq-shutdown taskq-enqueue taskq-dequeue taskq-complete taskq-resched)
  8. ;; task states
  9. (define STATE-QUEUED 0)
  10. (define STATE-EXECUTING 1)
  11. (define STATE-DONE 2)
  12. ;; message struct for task queue
  13. (struct qmsg [from type data] #:transparent)
  14. ;; sql to initialize db
  15. (define TASKQ-INIT "create table taskq (id blob primary key, task blob, notbefore int, state integer, completed int)")
  16. ;; runs the task queue
  17. (define (taskq-service c init-thread)
  18. ;; define some prepared statements
  19. (define delete-all-done (bind-prepared-statement
  20. (prepare c "delete from taskq where state=?")
  21. (list STATE-DONE)))
  22. (define delete-gc (prepare c "delete from taskq where state=? and completed<=?"))
  23. (define reset-executing (bind-prepared-statement
  24. (prepare c "update taskq set state=? where state=?")
  25. (list STATE-QUEUED STATE-EXECUTING)))
  26. (define create (prepare c "insert into taskq (id, task, notbefore, state, completed) values (?, ?, ?, ?, ?)"))
  27. (define update-state (prepare c "update taskq set state=? where id=?"))
  28. (define update-state-completed (prepare c "update taskq set state=?, completed=? where id=?"))
  29. (define update-notbefore (prepare c "update taskq set notbefore=?, state=? where id=?"))
  30. (define delete-by-id (prepare c "delete from taskq where id=?"))
  31. (define get-by-id (prepare c "select * from taskq where id=?"))
  32. (define get-next (prepare c "select * from taskq where state=? and notbefore<=? limit 1"))
  33. (define earliest-wakeup
  34. (bind-prepared-statement
  35. (prepare c "select notbefore from taskq where state=? order by notbefore asc limit 1")
  36. (list STATE-QUEUED)))
  37. ;; cleanup: if we crash and come back up, there might be done tasks and executing tasks
  38. ;; delete the done tasks, we don't have waiters for those anymore and any code that creates
  39. ;; waiters should check for the result it wanted and re-queue if needed
  40. ;; then, move executing tasks (which must have been interrupted by the crash) back to
  41. ;; queued so we can run them again
  42. (query-exec c delete-all-done)
  43. (query-exec c reset-executing)
  44. ;; done with init
  45. (thread-send init-thread #t #f)
  46. (set! init-thread #f)
  47. ;; waiters for a given task
  48. (define task-waiters (make-hash))
  49. ;; workers that are waiting for work
  50. (define dequeue-waiters (mutable-set))
  51. ;; adds a waiter for a task with given id
  52. (define (add-task-waiter! id from)
  53. (hash-update! task-waiters id (lambda (s) (set-add s from)) set))
  54. ;; invokes and removes all waiters for a given id
  55. (define (invoke-task-waiters! id)
  56. (for ([w (in-set (hash-ref task-waiters id set))])
  57. (thread-send w id #f))
  58. (hash-remove! task-waiters id))
  59. ;; enqueues data. wait? specifies whether the queuer should get a message back immediately with
  60. ;; the task id, or only upon completion
  61. (define (enqueue from data wait?)
  62. (match-define (list id-tmp task notbefore) data)
  63. (define id (or id-tmp (crypto-random-bytes 16)))
  64. (with-handlers ([exn:fail:sql? (lambda (e) (thread-send from e #f))])
  65. (match (query-maybe-row c get-by-id id)
  66. [#f
  67. (query-exec c create id (s-exp->fasl task) notbefore STATE-QUEUED 0)
  68. (if wait?
  69. (add-task-waiter! id from)
  70. (thread-send from id #f))]
  71. [(vector id task notbefore state completed)
  72. (if (and wait? (not (= state STATE-DONE)))
  73. (add-task-waiter! id from)
  74. (thread-send from id #f))]))
  75. (try-dequeue))
  76. ;; adds from to the dequeue list, attempts immediate dequeue if possible
  77. (define (dequeue from)
  78. (set-add! dequeue-waiters from)
  79. (try-dequeue))
  80. ;; errors all current waiters and shuts down the queue
  81. (define (shutdown)
  82. (for ([w (in-set dequeue-waiters)])
  83. (thread-send w #f #f))
  84. (for ([(k v) (in-hash task-waiters)])
  85. (for ([w (in-set v)])
  86. (thread-send w (error "queue shutdown") #f))))
  87. ;; attempts to dequeue the next work
  88. ;; if there is work, set it to executing and invoke and remove one of the workers
  89. ;; returns #t if there might be more work available, #f if there is guaranteed no work right now
  90. (define (try-dequeue)
  91. (if (set-empty? dequeue-waiters)
  92. #f
  93. (let ([w (set-first dequeue-waiters)])
  94. (with-handlers ([exn:fail:sql? (lambda (e) (displayln e) #f)])
  95. (match (query-maybe-row c get-next STATE-QUEUED (current-seconds))
  96. [#f #f]
  97. [(vector id task notbefore state completed)
  98. (query-exec c update-state STATE-EXECUTING id)
  99. (set-remove! dequeue-waiters w)
  100. (thread-send w (list id (fasl->s-exp task)) #f)
  101. #t])))))
  102. ;; marks a task as completed. invokes and removes any waiters
  103. (define (complete from id)
  104. (with-handlers ([exn:fail:sql? (lambda (e) (thread-send from e #f))])
  105. (query-exec c update-state-completed STATE-DONE (current-seconds) id)
  106. (invoke-task-waiters! id)
  107. (thread-send from #t #f)))
  108. ;; reschedules a task for a later time
  109. (define (resched from data)
  110. (match-define (list id notbefore) data)
  111. (with-handlers ([exn:fail:sql? (lambda (e) (thread-send from e #f))])
  112. (query-exec c update-notbefore notbefore STATE-QUEUED id)
  113. (thread-send from #t #f)))
  114. ;; collects old done tasks
  115. (define (collect-garbage now)
  116. (with-handlers ([exn:fail:sql? (lambda (e) (displayln e))])
  117. (query-exec c delete-gc STATE-DONE (- now 300))))
  118. ;; when was collect-garbage run last
  119. (define last-garbage (box (current-seconds)))
  120. ;; main task queue loop
  121. (let loop ()
  122. (define now (current-seconds))
  123. ;; collect garbage if needed
  124. (when (> now (+ (unbox last-garbage) 300))
  125. (collect-garbage)
  126. (set-box! last-garbage now))
  127. ;; get the next wakeup time
  128. (define wakeup (with-handlers ([exn:fail:sql? (lambda (e) 0)])
  129. (or (query-maybe-value c earliest-wakeup) +inf.0)))
  130. ;; get thread mail event
  131. (define recv-evt (thread-receive-evt))
  132. ;; create a wakeup event if wakeup is needed sometime
  133. (define wakeup-evt (if (>= wakeup now)
  134. (alarm-evt (* 1000 wakeup))
  135. always-evt))
  136. ;; wait for either a message or a wakeup
  137. ;; if both events are ready, one will be pseudorandomly chosen so it should be fair
  138. (define sync-result (sync recv-evt wakeup-evt))
  139. (cond
  140. [(equal? sync-result wakeup-evt)
  141. ;; run dequeue
  142. (let dequeue-loop ()
  143. (when (try-dequeue)
  144. (dequeue-loop)))
  145. (loop)]
  146. [else
  147. ;; handle message
  148. (match-define (qmsg from type data) (thread-receive))
  149. (match type
  150. ['shutdown (shutdown)]
  151. ['enqueue (enqueue from data #f) (loop)]
  152. ['enqueue-wait (enqueue from data #t) (loop)]
  153. ['dequeue (dequeue from) (loop)]
  154. ['complete (complete from data) (loop)]
  155. ['resched (resched from data) (loop)]
  156. [_ (thread-send from (error "unknown message") #f) (loop)])]))
  157. ;; disconnect and exit
  158. (disconnect c)
  159. (void))
  160. ;; creates a new taskq with db at given path, and optionally initializes it
  161. (define (make-taskq path [initialize? #f])
  162. (define c
  163. (if initialize?
  164. (sqlite3-connect #:database path #:mode 'create)
  165. (sqlite3-connect #:database path)))
  166. (when initialize?
  167. (query-exec c TASKQ-INIT))
  168. (define t (current-thread))
  169. (define q (thread (lambda () (taskq-service c t))))
  170. (thread-receive)
  171. q)
  172. ;; shuts down the task queue
  173. (define (taskq-shutdown q)
  174. (thread-send q (qmsg (current-thread) 'shutdown #f))
  175. (thread-wait q))
  176. ;; if the next mail is an exn, raise it, otherwise return the value
  177. (define (receive-check)
  178. (match (thread-receive)
  179. [(? exn? e) (raise e)]
  180. [result result]))
  181. ;; enqueues a task
  182. (define (taskq-enqueue q task [id #f] [notbefore 0] [wait? #t])
  183. (thread-send q (qmsg (current-thread) (if wait? 'enqueue-wait 'enqueue)
  184. (list id task notbefore)))
  185. (receive-check))
  186. ;; dequeues a task
  187. ;; returns the next task (now marked as executing) or #f
  188. (define (taskq-dequeue q)
  189. (thread-send q (qmsg (current-thread) 'dequeue #f))
  190. (receive-check))
  191. ;; marks a task complete
  192. (define (taskq-complete q id)
  193. (thread-send q (qmsg (current-thread) 'complete id))
  194. (receive-check))
  195. ;; reschedules a task for later
  196. ;; cannot be combined with taskq-complete: workers must call exactly one or the other, then call
  197. ;; taskq-dequeue for the next assignment
  198. (define (taskq-resched q id notbefore)
  199. (thread-send q (qmsg (current-thread) 'resched (list id notbefore)))
  200. (receive-check))