├── .gitignore ├── Cargo.lock ├── Cargo.toml ├── README.md └── src ├── error.rs ├── lib.rs ├── manager.rs ├── signal.rs └── test_util.rs /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | -------------------------------------------------------------------------------- /Cargo.lock: -------------------------------------------------------------------------------- 1 | # This file is automatically @generated by Cargo. 2 | # It is not intended for manual editing. 3 | version = 3 4 | 5 | [[package]] 6 | name = "autocfg" 7 | version = "1.1.0" 8 | source = "registry+https://github.com/rust-lang/crates.io-index" 9 | checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" 10 | 11 | [[package]] 12 | name = "bitflags" 13 | version = "1.3.2" 14 | source = "registry+https://github.com/rust-lang/crates.io-index" 15 | checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" 16 | 17 | [[package]] 18 | name = "bytes" 19 | version = "1.3.0" 20 | source = "registry+https://github.com/rust-lang/crates.io-index" 21 | checksum = "dfb24e866b15a1af2a1b663f10c6b6b8f397a84aadb828f12e5b289ec23a3a3c" 22 | 23 | [[package]] 24 | name = "cfg-if" 25 | version = "1.0.0" 26 | source = "registry+https://github.com/rust-lang/crates.io-index" 27 | checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" 28 | 29 | [[package]] 30 | name = "futures" 31 | version = "0.3.25" 32 | source = "registry+https://github.com/rust-lang/crates.io-index" 33 | checksum = "38390104763dc37a5145a53c29c63c1290b5d316d6086ec32c293f6736051bb0" 34 | dependencies = [ 35 | "futures-channel", 36 | "futures-core", 37 | "futures-executor", 38 | "futures-io", 39 | "futures-sink", 40 | "futures-task", 41 | "futures-util", 42 | ] 43 | 44 | [[package]] 45 | name = "futures-channel" 46 | version = "0.3.25" 47 | source = "registry+https://github.com/rust-lang/crates.io-index" 48 | checksum = "52ba265a92256105f45b719605a571ffe2d1f0fea3807304b522c1d778f79eed" 49 | dependencies = [ 50 | "futures-core", 51 | "futures-sink", 52 | ] 53 | 54 | [[package]] 55 | name = "futures-core" 56 | version = "0.3.25" 57 | source = "registry+https://github.com/rust-lang/crates.io-index" 58 | checksum = "04909a7a7e4633ae6c4a9ab280aeb86da1236243a77b694a49eacd659a4bd3ac" 59 | 60 | [[package]] 61 | name = "futures-executor" 62 | version = "0.3.25" 63 | source = "registry+https://github.com/rust-lang/crates.io-index" 64 | checksum = "7acc85df6714c176ab5edf386123fafe217be88c0840ec11f199441134a074e2" 65 | dependencies = [ 66 | "futures-core", 67 | "futures-task", 68 | "futures-util", 69 | ] 70 | 71 | [[package]] 72 | name = "futures-io" 73 | version = "0.3.25" 74 | source = "registry+https://github.com/rust-lang/crates.io-index" 75 | checksum = "00f5fb52a06bdcadeb54e8d3671f8888a39697dcb0b81b23b55174030427f4eb" 76 | 77 | [[package]] 78 | name = "futures-macro" 79 | version = "0.3.25" 80 | source = "registry+https://github.com/rust-lang/crates.io-index" 81 | checksum = "bdfb8ce053d86b91919aad980c220b1fb8401a9394410e1c289ed7e66b61835d" 82 | dependencies = [ 83 | "proc-macro2", 84 | "quote", 85 | "syn", 86 | ] 87 | 88 | [[package]] 89 | name = "futures-sink" 90 | version = "0.3.25" 91 | source = "registry+https://github.com/rust-lang/crates.io-index" 92 | checksum = "39c15cf1a4aa79df40f1bb462fb39676d0ad9e366c2a33b590d7c66f4f81fcf9" 93 | 94 | [[package]] 95 | name = "futures-task" 96 | version = "0.3.25" 97 | source = "registry+https://github.com/rust-lang/crates.io-index" 98 | checksum = "2ffb393ac5d9a6eaa9d3fdf37ae2776656b706e200c8e16b1bdb227f5198e6ea" 99 | 100 | [[package]] 101 | name = "futures-util" 102 | version = "0.3.25" 103 | source = "registry+https://github.com/rust-lang/crates.io-index" 104 | checksum = "197676987abd2f9cadff84926f410af1c183608d36641465df73ae8211dc65d6" 105 | dependencies = [ 106 | "futures-channel", 107 | "futures-core", 108 | "futures-io", 109 | "futures-macro", 110 | "futures-sink", 111 | "futures-task", 112 | "memchr", 113 | "pin-project-lite", 114 | "pin-utils", 115 | "slab", 116 | ] 117 | 118 | [[package]] 119 | name = "getrandom" 120 | version = "0.2.8" 121 | source = "registry+https://github.com/rust-lang/crates.io-index" 122 | checksum = "c05aeb6a22b8f62540c194aac980f2115af067bfe15a0734d7277a768d396b31" 123 | dependencies = [ 124 | "cfg-if", 125 | "libc", 126 | "wasi", 127 | ] 128 | 129 | [[package]] 130 | name = "hermit-abi" 131 | version = "0.2.6" 132 | source = "registry+https://github.com/rust-lang/crates.io-index" 133 | checksum = "ee512640fe35acbfb4bb779db6f0d80704c2cacfa2e39b601ef3e3f47d1ae4c7" 134 | dependencies = [ 135 | "libc", 136 | ] 137 | 138 | [[package]] 139 | name = "libc" 140 | version = "0.2.139" 141 | source = "registry+https://github.com/rust-lang/crates.io-index" 142 | checksum = "201de327520df007757c1f0adce6e827fe8562fbc28bfd9c15571c66ca1f5f79" 143 | 144 | [[package]] 145 | name = "lock_api" 146 | version = "0.4.9" 147 | source = "registry+https://github.com/rust-lang/crates.io-index" 148 | checksum = "435011366fe56583b16cf956f9df0095b405b82d76425bc8981c0e22e60ec4df" 149 | dependencies = [ 150 | "autocfg", 151 | "scopeguard", 152 | ] 153 | 154 | [[package]] 155 | name = "log" 156 | version = "0.4.17" 157 | source = "registry+https://github.com/rust-lang/crates.io-index" 158 | checksum = "abb12e687cfb44aa40f41fc3978ef76448f9b6038cad6aef4259d3c095a2382e" 159 | dependencies = [ 160 | "cfg-if", 161 | ] 162 | 163 | [[package]] 164 | name = "maplit" 165 | version = "1.0.2" 166 | source = "registry+https://github.com/rust-lang/crates.io-index" 167 | checksum = "3e2e65a1a2e43cfcb47a895c4c8b10d1f4a61097f9f254f183aee60cad9c651d" 168 | 169 | [[package]] 170 | name = "memchr" 171 | version = "2.5.0" 172 | source = "registry+https://github.com/rust-lang/crates.io-index" 173 | checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" 174 | 175 | [[package]] 176 | name = "mio" 177 | version = "0.8.5" 178 | source = "registry+https://github.com/rust-lang/crates.io-index" 179 | checksum = "e5d732bc30207a6423068df043e3d02e0735b155ad7ce1a6f76fe2baa5b158de" 180 | dependencies = [ 181 | "libc", 182 | "log", 183 | "wasi", 184 | "windows-sys", 185 | ] 186 | 187 | [[package]] 188 | name = "num_cpus" 189 | version = "1.15.0" 190 | source = "registry+https://github.com/rust-lang/crates.io-index" 191 | checksum = "0fac9e2da13b5eb447a6ce3d392f23a29d8694bff781bf03a16cd9ac8697593b" 192 | dependencies = [ 193 | "hermit-abi", 194 | "libc", 195 | ] 196 | 197 | [[package]] 198 | name = "once_cell" 199 | version = "1.17.0" 200 | source = "registry+https://github.com/rust-lang/crates.io-index" 201 | checksum = "6f61fba1741ea2b3d6a1e3178721804bb716a68a6aeba1149b5d52e3d464ea66" 202 | 203 | [[package]] 204 | name = "parking_lot" 205 | version = "0.12.1" 206 | source = "registry+https://github.com/rust-lang/crates.io-index" 207 | checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" 208 | dependencies = [ 209 | "lock_api", 210 | "parking_lot_core", 211 | ] 212 | 213 | [[package]] 214 | name = "parking_lot_core" 215 | version = "0.9.5" 216 | source = "registry+https://github.com/rust-lang/crates.io-index" 217 | checksum = "7ff9f3fef3968a3ec5945535ed654cb38ff72d7495a25619e2247fb15a2ed9ba" 218 | dependencies = [ 219 | "cfg-if", 220 | "libc", 221 | "redox_syscall", 222 | "smallvec", 223 | "windows-sys", 224 | ] 225 | 226 | [[package]] 227 | name = "pin-project-lite" 228 | version = "0.2.9" 229 | source = "registry+https://github.com/rust-lang/crates.io-index" 230 | checksum = "e0a7ae3ac2f1173085d398531c705756c94a4c56843785df85a60c1a0afac116" 231 | 232 | [[package]] 233 | name = "pin-utils" 234 | version = "0.1.0" 235 | source = "registry+https://github.com/rust-lang/crates.io-index" 236 | checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" 237 | 238 | [[package]] 239 | name = "ppv-lite86" 240 | version = "0.2.17" 241 | source = "registry+https://github.com/rust-lang/crates.io-index" 242 | checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" 243 | 244 | [[package]] 245 | name = "proc-macro2" 246 | version = "1.0.49" 247 | source = "registry+https://github.com/rust-lang/crates.io-index" 248 | checksum = "57a8eca9f9c4ffde41714334dee777596264c7825420f521abc92b5b5deb63a5" 249 | dependencies = [ 250 | "unicode-ident", 251 | ] 252 | 253 | [[package]] 254 | name = "quote" 255 | version = "1.0.23" 256 | source = "registry+https://github.com/rust-lang/crates.io-index" 257 | checksum = "8856d8364d252a14d474036ea1358d63c9e6965c8e5c1885c18f73d70bff9c7b" 258 | dependencies = [ 259 | "proc-macro2", 260 | ] 261 | 262 | [[package]] 263 | name = "rand" 264 | version = "0.8.5" 265 | source = "registry+https://github.com/rust-lang/crates.io-index" 266 | checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" 267 | dependencies = [ 268 | "libc", 269 | "rand_chacha", 270 | "rand_core", 271 | ] 272 | 273 | [[package]] 274 | name = "rand_chacha" 275 | version = "0.3.1" 276 | source = "registry+https://github.com/rust-lang/crates.io-index" 277 | checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" 278 | dependencies = [ 279 | "ppv-lite86", 280 | "rand_core", 281 | ] 282 | 283 | [[package]] 284 | name = "rand_core" 285 | version = "0.6.4" 286 | source = "registry+https://github.com/rust-lang/crates.io-index" 287 | checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" 288 | dependencies = [ 289 | "getrandom", 290 | ] 291 | 292 | [[package]] 293 | name = "redox_syscall" 294 | version = "0.2.16" 295 | source = "registry+https://github.com/rust-lang/crates.io-index" 296 | checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" 297 | dependencies = [ 298 | "bitflags", 299 | ] 300 | 301 | [[package]] 302 | name = "scopeguard" 303 | version = "1.1.0" 304 | source = "registry+https://github.com/rust-lang/crates.io-index" 305 | checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" 306 | 307 | [[package]] 308 | name = "signal-hook-registry" 309 | version = "1.4.0" 310 | source = "registry+https://github.com/rust-lang/crates.io-index" 311 | checksum = "e51e73328dc4ac0c7ccbda3a494dfa03df1de2f46018127f60c693f2648455b0" 312 | dependencies = [ 313 | "libc", 314 | ] 315 | 316 | [[package]] 317 | name = "slab" 318 | version = "0.4.7" 319 | source = "registry+https://github.com/rust-lang/crates.io-index" 320 | checksum = "4614a76b2a8be0058caa9dbbaf66d988527d86d003c11a94fbd335d7661edcef" 321 | dependencies = [ 322 | "autocfg", 323 | ] 324 | 325 | [[package]] 326 | name = "smallvec" 327 | version = "1.10.0" 328 | source = "registry+https://github.com/rust-lang/crates.io-index" 329 | checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0" 330 | 331 | [[package]] 332 | name = "socket2" 333 | version = "0.4.7" 334 | source = "registry+https://github.com/rust-lang/crates.io-index" 335 | checksum = "02e2d2db9033d13a1567121ddd7a095ee144db4e1ca1b1bda3419bc0da294ebd" 336 | dependencies = [ 337 | "libc", 338 | "winapi", 339 | ] 340 | 341 | [[package]] 342 | name = "syn" 343 | version = "1.0.107" 344 | source = "registry+https://github.com/rust-lang/crates.io-index" 345 | checksum = "1f4064b5b16e03ae50984a5a8ed5d4f8803e6bc1fd170a3cda91a1be4b18e3f5" 346 | dependencies = [ 347 | "proc-macro2", 348 | "quote", 349 | "unicode-ident", 350 | ] 351 | 352 | [[package]] 353 | name = "task-motel" 354 | version = "0.1.0" 355 | dependencies = [ 356 | "futures", 357 | "maplit", 358 | "parking_lot", 359 | "rand", 360 | "tokio", 361 | "tokio-stream", 362 | "tracing", 363 | ] 364 | 365 | [[package]] 366 | name = "tokio" 367 | version = "1.23.1" 368 | source = "registry+https://github.com/rust-lang/crates.io-index" 369 | checksum = "38a54aca0c15d014013256222ba0ebed095673f89345dd79119d912eb561b7a8" 370 | dependencies = [ 371 | "autocfg", 372 | "bytes", 373 | "libc", 374 | "memchr", 375 | "mio", 376 | "num_cpus", 377 | "parking_lot", 378 | "pin-project-lite", 379 | "signal-hook-registry", 380 | "socket2", 381 | "tokio-macros", 382 | "windows-sys", 383 | ] 384 | 385 | [[package]] 386 | name = "tokio-macros" 387 | version = "1.8.2" 388 | source = "registry+https://github.com/rust-lang/crates.io-index" 389 | checksum = "d266c00fde287f55d3f1c3e96c500c362a2b8c695076ec180f27918820bc6df8" 390 | dependencies = [ 391 | "proc-macro2", 392 | "quote", 393 | "syn", 394 | ] 395 | 396 | [[package]] 397 | name = "tokio-stream" 398 | version = "0.1.11" 399 | source = "registry+https://github.com/rust-lang/crates.io-index" 400 | checksum = "d660770404473ccd7bc9f8b28494a811bc18542b915c0855c51e8f419d5223ce" 401 | dependencies = [ 402 | "futures-core", 403 | "pin-project-lite", 404 | "tokio", 405 | "tokio-util", 406 | ] 407 | 408 | [[package]] 409 | name = "tokio-util" 410 | version = "0.7.4" 411 | source = "registry+https://github.com/rust-lang/crates.io-index" 412 | checksum = "0bb2e075f03b3d66d8d8785356224ba688d2906a371015e225beeb65ca92c740" 413 | dependencies = [ 414 | "bytes", 415 | "futures-core", 416 | "futures-sink", 417 | "pin-project-lite", 418 | "tokio", 419 | ] 420 | 421 | [[package]] 422 | name = "tracing" 423 | version = "0.1.37" 424 | source = "registry+https://github.com/rust-lang/crates.io-index" 425 | checksum = "8ce8c33a8d48bd45d624a6e523445fd21ec13d3653cd51f681abf67418f54eb8" 426 | dependencies = [ 427 | "cfg-if", 428 | "pin-project-lite", 429 | "tracing-attributes", 430 | "tracing-core", 431 | ] 432 | 433 | [[package]] 434 | name = "tracing-attributes" 435 | version = "0.1.23" 436 | source = "registry+https://github.com/rust-lang/crates.io-index" 437 | checksum = "4017f8f45139870ca7e672686113917c71c7a6e02d4924eda67186083c03081a" 438 | dependencies = [ 439 | "proc-macro2", 440 | "quote", 441 | "syn", 442 | ] 443 | 444 | [[package]] 445 | name = "tracing-core" 446 | version = "0.1.30" 447 | source = "registry+https://github.com/rust-lang/crates.io-index" 448 | checksum = "24eb03ba0eab1fd845050058ce5e616558e8f8d8fca633e6b163fe25c797213a" 449 | dependencies = [ 450 | "once_cell", 451 | ] 452 | 453 | [[package]] 454 | name = "unicode-ident" 455 | version = "1.0.6" 456 | source = "registry+https://github.com/rust-lang/crates.io-index" 457 | checksum = "84a22b9f218b40614adcb3f4ff08b703773ad44fa9423e4e0d346d5db86e4ebc" 458 | 459 | [[package]] 460 | name = "wasi" 461 | version = "0.11.0+wasi-snapshot-preview1" 462 | source = "registry+https://github.com/rust-lang/crates.io-index" 463 | checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" 464 | 465 | [[package]] 466 | name = "winapi" 467 | version = "0.3.9" 468 | source = "registry+https://github.com/rust-lang/crates.io-index" 469 | checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" 470 | dependencies = [ 471 | "winapi-i686-pc-windows-gnu", 472 | "winapi-x86_64-pc-windows-gnu", 473 | ] 474 | 475 | [[package]] 476 | name = "winapi-i686-pc-windows-gnu" 477 | version = "0.4.0" 478 | source = "registry+https://github.com/rust-lang/crates.io-index" 479 | checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" 480 | 481 | [[package]] 482 | name = "winapi-x86_64-pc-windows-gnu" 483 | version = "0.4.0" 484 | source = "registry+https://github.com/rust-lang/crates.io-index" 485 | checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" 486 | 487 | [[package]] 488 | name = "windows-sys" 489 | version = "0.42.0" 490 | source = "registry+https://github.com/rust-lang/crates.io-index" 491 | checksum = "5a3e1820f08b8513f676f7ab6c1f99ff312fb97b553d30ff4dd86f9f15728aa7" 492 | dependencies = [ 493 | "windows_aarch64_gnullvm", 494 | "windows_aarch64_msvc", 495 | "windows_i686_gnu", 496 | "windows_i686_msvc", 497 | "windows_x86_64_gnu", 498 | "windows_x86_64_gnullvm", 499 | "windows_x86_64_msvc", 500 | ] 501 | 502 | [[package]] 503 | name = "windows_aarch64_gnullvm" 504 | version = "0.42.0" 505 | source = "registry+https://github.com/rust-lang/crates.io-index" 506 | checksum = "41d2aa71f6f0cbe00ae5167d90ef3cfe66527d6f613ca78ac8024c3ccab9a19e" 507 | 508 | [[package]] 509 | name = "windows_aarch64_msvc" 510 | version = "0.42.0" 511 | source = "registry+https://github.com/rust-lang/crates.io-index" 512 | checksum = "dd0f252f5a35cac83d6311b2e795981f5ee6e67eb1f9a7f64eb4500fbc4dcdb4" 513 | 514 | [[package]] 515 | name = "windows_i686_gnu" 516 | version = "0.42.0" 517 | source = "registry+https://github.com/rust-lang/crates.io-index" 518 | checksum = "fbeae19f6716841636c28d695375df17562ca208b2b7d0dc47635a50ae6c5de7" 519 | 520 | [[package]] 521 | name = "windows_i686_msvc" 522 | version = "0.42.0" 523 | source = "registry+https://github.com/rust-lang/crates.io-index" 524 | checksum = "84c12f65daa39dd2babe6e442988fc329d6243fdce47d7d2d155b8d874862246" 525 | 526 | [[package]] 527 | name = "windows_x86_64_gnu" 528 | version = "0.42.0" 529 | source = "registry+https://github.com/rust-lang/crates.io-index" 530 | checksum = "bf7b1b21b5362cbc318f686150e5bcea75ecedc74dd157d874d754a2ca44b0ed" 531 | 532 | [[package]] 533 | name = "windows_x86_64_gnullvm" 534 | version = "0.42.0" 535 | source = "registry+https://github.com/rust-lang/crates.io-index" 536 | checksum = "09d525d2ba30eeb3297665bd434a54297e4170c7f1a44cad4ef58095b4cd2028" 537 | 538 | [[package]] 539 | name = "windows_x86_64_msvc" 540 | version = "0.42.0" 541 | source = "registry+https://github.com/rust-lang/crates.io-index" 542 | checksum = "f40009d85759725a34da6d89a94e63d7bdc50a862acf0dbc7c8e488f1edcb6f5" 543 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "task-motel" 3 | version = "0.1.0" 4 | authors = ["Michael Dougherty "] 5 | edition = "2021" 6 | description = "Opinionated (Tokio) task manager with nested task groups and stoppable tasks" 7 | license = "Apache-2.0" 8 | homepage = "https://github.com/maackle/task-motel-rs" 9 | documentation = "https://docs.rs/task-motel" 10 | 11 | [dependencies] 12 | futures = "0.3" 13 | parking_lot = "0.12" 14 | tokio = { version = "1.11", features = [ "full" ] } 15 | tracing = "0.1" 16 | 17 | [dev-dependencies] 18 | maplit = "1.0" 19 | rand = "0.8" 20 | tokio = { version = "1.11", features = [ "full", "test-util" ] } 21 | tokio-stream = { version = "0.1", features = [ "sync" ] } 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # task-motel 2 | 3 | An opinionated Tokio* task manager with arbitrary nested groupings of tasks. 4 | 5 | The task manager takes two generic parameters: 6 | 7 | ```rust 8 | pub struct TaskManager 9 | where 10 | GroupKey: Clone + Eq + Hash + Send + std::fmt::Debug + 'static, 11 | Outcome: Send + 'static, 12 | ``` 13 | 14 | Each distinct `GroupKey` (as defined by `Eq`) corresponds to a distinct group, so tasks registered with the same `GroupKey` are in the same group. It is useful to use an `enum` for your `GroupKey`. 15 | 16 | The `Outcome` is the type that all tasks must return, which may be whatever you want. 17 | 18 | \* task-motel is built to be almost entirely executor-agnostic, but ultimately I needed the ability to spawn a task which would run to completion without polling, which I could not find without leaning on some executor or another. If this can be solved, then the Tokio dependency can be removed. 19 | 20 | ## Handling task outcomes 21 | When constructing a `TaskManager`, you are given a channel receiver, which receives `(GroupKey, Outcome)` for each task that has completed. Depending on the outcome of a task, you may respond by terminating the application, logging an error, spawning new tasks, or whatever else you want. 22 | 23 | ## Grouped tasks 24 | When registering a task, you assign it to a group. That group of tasks can be managed separately from other groups of tasks: the entire group can be instructed to stop, and the shutdown can be awaited. 25 | 26 | ## Nested groups 27 | Groups can be arbitrarily nested, so that shutting down a parent group will also shut down all child groups recursively. The nesting structure is determined by `Fn(GroupKey) -> Option` function argument to the `TaskManager::new` constructor, which takes a `GroupKey` and produces an `Option` declaring the parent of the group, defining the tree structure. 28 | 29 | ## Stopping tasks 30 | When registering a task, you must do so with a closure that gets a `StopListener` injected into it. This `StopListener` is a simple future which is resolved when the task manager receives a shutdown signal, so it should be incorporated into your task accordingly, to allow the task to shut down from external input. 31 | 32 | The StopListener should be dropped only after the task is complete, because the task manager uses the number of StopListeners in circulation to determine how many tasks are being actively managed. 33 | 34 | -------------------------------------------------------------------------------- /src/error.rs: -------------------------------------------------------------------------------- 1 | #![allow(missing_docs)] 2 | 3 | pub type TmResult = Result; 4 | 5 | /// An error that is thrown from within a managed task 6 | pub trait TaskError { 7 | fn is_recoverable(&self) -> bool; 8 | } 9 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | // #![warn(missing_docs)] 2 | 3 | mod error; 4 | pub use error::*; 5 | 6 | mod manager; 7 | use futures::task::FutureObj; 8 | pub use manager::*; 9 | 10 | mod signal; 11 | pub use signal::*; 12 | 13 | #[cfg(test)] 14 | pub mod test_util; 15 | 16 | /// A JoinHandle returning the result of running the task 17 | pub type Task = FutureObj<'static, ()>; 18 | -------------------------------------------------------------------------------- /src/manager.rs: -------------------------------------------------------------------------------- 1 | //! Manage tasks arranged in nested groups 2 | //! 3 | //! Groups can be added and removed dynamically. When a group is removed, 4 | //! all of its tasks are stopped, and all of its descendent groups are also removed, 5 | //! and their contained tasks stopped as well. The group is only completely removed when 6 | //! all descendent tasks have stopped. 7 | 8 | use std::{ 9 | collections::{HashMap, HashSet}, 10 | hash::Hash, 11 | sync::{atomic::AtomicU32, Arc}, 12 | }; 13 | 14 | use futures::{ 15 | channel::mpsc, future::BoxFuture, stream::FuturesUnordered, Future, FutureExt, StreamExt, 16 | }; 17 | use tracing::Instrument; 18 | 19 | use crate::{signal::StopListener, StopBroadcaster}; 20 | 21 | /// Tracks tasks at the global conductor level, as well as each individual cell level. 22 | pub struct TaskManager { 23 | span: Option, 24 | groups: HashMap, 25 | children: HashMap>, 26 | parent_map: Box Option>, 27 | outcome_rx: mpsc::Sender<(GroupKey, Outcome)>, 28 | // used to keep track of the number of tasks still running 29 | stopping_group_counts: Vec>, 30 | } 31 | 32 | impl TaskManager 33 | where 34 | GroupKey: Clone + Eq + Hash + Send + std::fmt::Debug + 'static, 35 | Outcome: Send + 'static, 36 | { 37 | pub fn new( 38 | outcome_rx: mpsc::Sender<(GroupKey, Outcome)>, 39 | parent_map: impl 'static + Send + Sync + Fn(&GroupKey) -> Option + 'static, 40 | ) -> Self { 41 | Self { 42 | span: None, 43 | groups: Default::default(), 44 | children: Default::default(), 45 | parent_map: Box::new(parent_map), 46 | outcome_rx, 47 | stopping_group_counts: Default::default(), 48 | } 49 | } 50 | 51 | pub fn new_instrumented( 52 | span: tracing::Span, 53 | outcome_rx: mpsc::Sender<(GroupKey, Outcome)>, 54 | parent_map: impl 'static + Send + Sync + Fn(&GroupKey) -> Option + 'static, 55 | ) -> Self { 56 | Self { 57 | span: Some(span), 58 | groups: Default::default(), 59 | children: Default::default(), 60 | parent_map: Box::new(parent_map), 61 | outcome_rx, 62 | stopping_group_counts: Default::default(), 63 | } 64 | } 65 | 66 | /// Add a task to a group 67 | pub fn add_task + Send + 'static>( 68 | &mut self, 69 | key: GroupKey, 70 | f: impl FnOnce(StopListener) -> Fut + Send + 'static, 71 | ) { 72 | let span = self.span.clone(); 73 | let mut tx = self.outcome_rx.clone(); 74 | let group = self.group(key.clone()); 75 | let listener = group.stopper.listener(); 76 | let task = async move { 77 | let outcome = if let Some(span) = span { 78 | f(listener).instrument(span).await 79 | } else { 80 | f(listener).await 81 | }; 82 | tx.try_send((key, outcome)).ok(); 83 | } 84 | .boxed(); 85 | group.tasks.spawn(task); 86 | } 87 | 88 | /// Remove a group, returning a future that can be waited upon for all tasks 89 | /// in that group to complete. 90 | pub fn stop_group(&mut self, key: &GroupKey) -> GroupStop { 91 | let mut js = tokio::task::JoinSet::new(); 92 | for key in self.descendants(key) { 93 | if let Some(mut group) = self.groups.remove(&key) { 94 | // Signal all tasks to stop. 95 | group.stopper.emit(); 96 | let num = group.stopper.num; 97 | self.stopping_group_counts.push(num); 98 | 99 | js.spawn(finish_joinset(group.tasks)); 100 | } 101 | } 102 | 103 | async move { finish_joinset(js).await }.boxed() 104 | } 105 | 106 | pub(crate) fn descendants(&self, key: &GroupKey) -> HashSet { 107 | let mut all = HashSet::new(); 108 | all.insert(key.clone()); 109 | 110 | let this = &self; 111 | 112 | if let Some(children) = this.children.get(&key) { 113 | for child in children { 114 | all.extend(this.descendants(child)); 115 | } 116 | } 117 | 118 | all 119 | } 120 | 121 | fn group(&mut self, key: GroupKey) -> &mut TaskGroup { 122 | self.groups.entry(key.clone()).or_insert_with(|| { 123 | if let Some(parent) = (self.parent_map)(&key) { 124 | self.children 125 | .entry(parent) 126 | .or_insert_with(HashSet::new) 127 | .insert(key); 128 | } 129 | TaskGroup::new() 130 | }) 131 | } 132 | 133 | /// For testing purposes only, this is not a reliable indicator, since the JoinSet 134 | /// is not polled except when a group is ending, so the count never actually 135 | /// decreases. 136 | #[cfg(test)] 137 | fn num_tasks(&self, key: &GroupKey) -> usize { 138 | let current = self 139 | .groups 140 | .get(key) 141 | .map(|group| group.tasks.len()) 142 | .unwrap_or_default(); 143 | 144 | let pending = self 145 | .stopping_group_counts 146 | .iter() 147 | .map(|c| c.load(std::sync::atomic::Ordering::SeqCst)) 148 | .sum::() as usize; 149 | 150 | // dbg!(current) + dbg!(pending) 151 | current + pending 152 | } 153 | } 154 | 155 | pub type GroupStop = BoxFuture<'static, ()>; 156 | 157 | struct TaskGroup { 158 | pub(crate) tasks: tokio::task::JoinSet<()>, 159 | pub(crate) stopper: StopBroadcaster, 160 | } 161 | 162 | impl TaskGroup { 163 | pub fn new() -> Self { 164 | Self { 165 | tasks: tokio::task::JoinSet::new(), 166 | stopper: StopBroadcaster::new(), 167 | } 168 | } 169 | } 170 | 171 | pub type TaskStream = 172 | futures::stream::SelectAll>>; 173 | 174 | async fn finish_joinset(mut js: tokio::task::JoinSet<()>) { 175 | futures::stream::unfold(&mut js, |tasks| async move { 176 | if let Err(err) = tasks.join_next().await? { 177 | tracing::error!("task_motel: Error while joining task: {:?}", err); 178 | } 179 | Some(((), tasks)) 180 | }) 181 | .collect::>() 182 | .await; 183 | js.detach_all(); 184 | } 185 | #[cfg(test)] 186 | mod tests { 187 | use futures::{channel::mpsc, SinkExt}; 188 | use maplit::hashset; 189 | use rand::seq::SliceRandom; 190 | 191 | use crate::test_util::*; 192 | 193 | use super::*; 194 | 195 | #[derive(Debug, Clone, Hash, PartialEq, Eq)] 196 | enum GroupKey { 197 | A, 198 | B, 199 | C, 200 | D, 201 | E, 202 | F, 203 | G, 204 | } 205 | 206 | #[tokio::test(start_paused = true)] 207 | async fn test_task_completion() { 208 | use GroupKey::*; 209 | let (outcome_tx, mut outcome_rx) = mpsc::channel(1); 210 | let mut tm: TaskManager = TaskManager::new(outcome_tx, |g| match g { 211 | B => Some(A), 212 | _ => None, 213 | }); 214 | 215 | let sec = tokio::time::Duration::from_secs(1); 216 | 217 | tm.add_task(A, move |stop| { 218 | async move { 219 | let _stop = stop; 220 | tokio::time::sleep(sec).await; 221 | tokio::time::sleep(sec).await; 222 | tokio::time::sleep(sec).await; 223 | "done".to_string() 224 | } 225 | .boxed() 226 | }); 227 | 228 | tokio::time::advance(sec).await; 229 | 230 | assert_eq!(tm.num_tasks(&A), 1); 231 | 232 | tokio::time::advance(sec).await; 233 | 234 | let stopping = tm.stop_group(&A); 235 | 236 | assert_eq!(tm.num_tasks(&A), 1); 237 | 238 | stopping.await; 239 | 240 | assert_eq!(tm.num_tasks(&A), 0); 241 | 242 | // tm.stop_group(&A).await; 243 | assert_eq!(outcome_rx.next().await.unwrap(), (A, "done".to_string())); 244 | assert_eq!(tm.num_tasks(&A), 0); 245 | } 246 | 247 | #[tokio::test] 248 | async fn test_descendants() { 249 | use GroupKey::*; 250 | let (outcome_tx, outcome_rx) = mpsc::channel(1); 251 | let mut tm: TaskManager = TaskManager::new(outcome_tx, |g| match g { 252 | A => None, 253 | B => Some(A), 254 | C => Some(B), 255 | D => Some(B), 256 | E => Some(D), 257 | F => Some(E), 258 | G => Some(C), 259 | }); 260 | 261 | let mut keys = vec![A, B, C, D, E, F, G]; 262 | keys.shuffle(&mut rand::thread_rng()); 263 | 264 | // Set up the parent map in random order 265 | for key in keys.clone() { 266 | tm.add_task(key.clone(), |_| async move { format!("{:?}", key) }) 267 | } 268 | 269 | assert_eq!(tm.descendants(&A), hashset! {A, B, C, D, E, F, G}); 270 | assert_eq!(tm.descendants(&B), hashset! {B, C, D, E, F, G}); 271 | assert_eq!(tm.descendants(&C), hashset! {C, G}); 272 | assert_eq!(tm.descendants(&D), hashset! {D, E, F}); 273 | assert_eq!(tm.descendants(&E), hashset! {E, F}); 274 | assert_eq!(tm.descendants(&F), hashset! {F}); 275 | assert_eq!(tm.descendants(&G), hashset! {G}); 276 | 277 | tm.stop_group(&A).await; 278 | 279 | assert_eq!( 280 | outcome_rx.take(keys.len()).collect::>().await, 281 | hashset! { 282 | (A, "A".to_string()), 283 | (B, "B".to_string()), 284 | (C, "C".to_string()), 285 | (D, "D".to_string()), 286 | (E, "E".to_string()), 287 | (F, "F".to_string()), 288 | (G, "G".to_string()), 289 | } 290 | ); 291 | } 292 | 293 | #[tokio::test] 294 | async fn test_group_nesting() { 295 | use GroupKey::*; 296 | let (outcome_tx, mut outcome_rx) = mpsc::channel(1); 297 | let (mut trigger_tx, trigger_rx) = mpsc::channel(1); 298 | let mut tm: TaskManager = TaskManager::new(outcome_tx, |g| match g { 299 | A => None, 300 | B => Some(A), 301 | C => Some(B), 302 | D => Some(B), 303 | _ => None, 304 | }); 305 | 306 | tm.add_task(A, |stop| blocker("a1", stop)); 307 | tm.add_task(A, |stop| blocker("a2", stop)); 308 | tm.add_task(B, |stop| blocker("b1", stop)); 309 | tm.add_task(C, |stop| blocker("c1", stop)); 310 | tm.add_task(D, |stop| blocker("d1", stop)); 311 | tm.add_task(E, |stop| fused("e1", stop.fuse_with(trigger_rx.take(1)))); 312 | 313 | assert_eq!(tm.num_tasks(&A), 2); 314 | assert_eq!(tm.num_tasks(&B), 1); 315 | assert_eq!(tm.num_tasks(&C), 1); 316 | assert_eq!(tm.num_tasks(&D), 1); 317 | assert_eq!(tm.num_tasks(&E), 1); 318 | 319 | trigger_tx.send(()).await.unwrap(); 320 | assert_eq!(outcome_rx.next().await.unwrap(), (E, "e1".to_string())); 321 | // The actual task number will not decrease until the group is officially stopped. 322 | assert_eq!(tm.num_tasks(&E), 1); 323 | 324 | let stopping = tm.stop_group(&D); 325 | assert_eq!(tm.num_tasks(&D), 1); 326 | stopping.await; 327 | assert_eq!(tm.num_tasks(&D), 0); 328 | assert_eq!( 329 | hashset![outcome_rx.next().await.unwrap(),], 330 | hashset![(D, "d1".to_string())] 331 | ); 332 | 333 | assert_eq!(tm.num_tasks(&A), 2); 334 | assert_eq!(tm.num_tasks(&B), 1); 335 | assert_eq!(tm.num_tasks(&C), 1); 336 | assert_eq!(tm.num_tasks(&D), 0); 337 | 338 | tm.add_task(D, |stop| blocker("dx", stop)); 339 | assert_eq!(tm.num_tasks(&D), 1); 340 | 341 | tm.stop_group(&B).await; 342 | assert_eq!( 343 | hashset![ 344 | outcome_rx.next().await.unwrap(), 345 | outcome_rx.next().await.unwrap(), 346 | outcome_rx.next().await.unwrap(), 347 | ], 348 | hashset![ 349 | (B, "b1".to_string()), 350 | (C, "c1".to_string()), 351 | (D, "dx".to_string()) 352 | ] 353 | ); 354 | 355 | assert_eq!(tm.num_tasks(&A), 2); 356 | assert_eq!(tm.num_tasks(&B), 0); 357 | assert_eq!(tm.num_tasks(&C), 0); 358 | assert_eq!(tm.num_tasks(&D), 0); 359 | 360 | tm.add_task(D, |stop| blocker("dy", stop)); 361 | assert_eq!(tm.num_tasks(&D), 1); 362 | 363 | tm.stop_group(&A).await; 364 | assert_eq!( 365 | hashset![ 366 | outcome_rx.next().await.unwrap(), 367 | outcome_rx.next().await.unwrap(), 368 | outcome_rx.next().await.unwrap(), 369 | ], 370 | hashset![ 371 | (A, "a1".to_string()), 372 | (A, "a2".to_string()), 373 | (D, "dy".to_string()) 374 | ] 375 | ); 376 | } 377 | } 378 | -------------------------------------------------------------------------------- /src/signal.rs: -------------------------------------------------------------------------------- 1 | //! Stop signal broadcasters and receivers 2 | 3 | use std::{ 4 | pin::Pin, 5 | sync::{ 6 | atomic::{AtomicU32, Ordering}, 7 | Arc, 8 | }, 9 | task::{Context, Poll}, 10 | }; 11 | 12 | use futures::{channel::oneshot, stream::Fuse, Future, FutureExt, Stream, StreamExt}; 13 | use parking_lot::Mutex; 14 | use tokio::sync::Notify; 15 | 16 | #[derive(Clone)] 17 | pub struct StopBroadcaster { 18 | txs: Arc>>>, 19 | pub(crate) num: Arc, 20 | notify: Arc, 21 | } 22 | 23 | impl StopBroadcaster { 24 | pub fn new() -> Self { 25 | Self { 26 | txs: Arc::new(Mutex::new(vec![])), 27 | num: Arc::new(AtomicU32::new(0)), 28 | notify: Arc::new(Notify::new()), 29 | } 30 | } 31 | 32 | pub fn listener(&self) -> StopListener { 33 | self.num.fetch_add(1, Ordering::SeqCst); 34 | let notify = self.notify.clone(); 35 | let (tx, rx) = oneshot::channel(); 36 | 37 | self.txs.lock().push(tx); 38 | 39 | StopListener { 40 | receiver: rx, 41 | num: self.num.clone(), 42 | notify, 43 | } 44 | } 45 | 46 | pub fn emit(&mut self) { 47 | // If a receiver is dropped, we don't care. 48 | for tx in self.txs.lock().drain(..) { 49 | tx.send(()).ok(); 50 | } 51 | } 52 | 53 | pub fn len(&self) -> u32 { 54 | self.num.load(Ordering::SeqCst) 55 | } 56 | 57 | pub async fn until_empty(&self) { 58 | while self.len() > 0 { 59 | self.notify.notified().await 60 | } 61 | } 62 | } 63 | 64 | /// StopListener should be incorporated into each user-defined task. 65 | /// It Derefs to a channel receiver which can be awaited. When resolved, 66 | /// the task should shut itself down. 67 | /// 68 | /// When the StopListener is dropped, that signals the TaskManager that 69 | /// the task has ended. 70 | pub struct StopListener { 71 | receiver: oneshot::Receiver<()>, 72 | num: Arc, 73 | notify: Arc, 74 | } 75 | 76 | impl StopListener { 77 | /// Modify a stream so that when the StopListener resolves, the stream is fused 78 | /// and ends. 79 | pub fn fuse_with>( 80 | self, 81 | stream: S, 82 | ) -> Fuse>>> { 83 | StreamExt::fuse(Box::pin(StopListenerFuse { stream, stop: self })) 84 | } 85 | 86 | /// Accessor for the underlying oneshot receiver 87 | pub fn receiver(&mut self) -> &mut oneshot::Receiver<()> { 88 | &mut self.receiver 89 | } 90 | 91 | /// Returns false if the oneshot message has not yet been sent. 92 | /// Returns true if the message was sent or the channel canceled/dropped. 93 | pub fn ready(&mut self) -> bool { 94 | !matches!(self.receiver.try_recv(), Ok(None)) 95 | } 96 | } 97 | 98 | impl Drop for StopListener { 99 | fn drop(&mut self) { 100 | self.num.fetch_sub(1, Ordering::SeqCst); 101 | self.notify.notify_one(); 102 | } 103 | } 104 | 105 | impl Future for StopListener { 106 | type Output = (); 107 | 108 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 109 | match Box::pin(&mut self.receiver).poll_unpin(cx) { 110 | Poll::Ready(_) => Poll::Ready(()), 111 | Poll::Pending => Poll::Pending, 112 | } 113 | } 114 | } 115 | 116 | pub struct StopListenerFuse> { 117 | stream: S, 118 | stop: StopListener, 119 | } 120 | 121 | impl> Stream for StopListenerFuse { 122 | type Item = T; 123 | 124 | fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 125 | if let Poll::Ready(_) = Box::pin(&mut self.stop).poll_unpin(cx) { 126 | return Poll::Ready(None); 127 | } 128 | 129 | Stream::poll_next(Pin::new(&mut self.stream), cx) 130 | } 131 | } 132 | 133 | #[cfg(test)] 134 | mod tests { 135 | use super::*; 136 | use crate::test_util::*; 137 | 138 | #[tokio::test] 139 | async fn test_stop_empty() { 140 | let x = StopBroadcaster::new(); 141 | assert_eq!(x.len(), 0); 142 | assert!(ready(x.until_empty()).await); 143 | } 144 | 145 | #[tokio::test] 146 | async fn test_stop() { 147 | let mut x = StopBroadcaster::new(); 148 | let a = x.listener(); 149 | let mut b = x.listener(); 150 | let c = x.listener(); 151 | assert_eq!(x.len(), 3); 152 | assert!(not_ready(x.until_empty()).await); 153 | 154 | assert!(not_ready(a).await); 155 | assert_eq!(x.len(), 2); 156 | assert!(!b.ready()); 157 | 158 | x.emit(); 159 | assert!(b.ready()); 160 | assert!(ready(b).await); 161 | assert_eq!(x.len(), 1); 162 | assert!(not_ready(x.until_empty()).await); 163 | 164 | assert!(ready(c).await); 165 | assert_eq!(x.len(), 0); 166 | assert!(ready(x.until_empty()).await); 167 | 168 | let y = StopBroadcaster::new(); 169 | let mut d = y.listener(); 170 | drop(y); 171 | assert!(d.ready()); 172 | assert!(ready(d).await); 173 | } 174 | 175 | #[tokio::test] 176 | async fn test_fuse_with() { 177 | { 178 | let mut tx = StopBroadcaster::new(); 179 | let rx = tx.listener(); 180 | let mut fused = rx.fuse_with(futures::stream::repeat(0)); 181 | assert_eq!(fused.next().await, Some(0)); 182 | assert_eq!(fused.next().await, Some(0)); 183 | tx.emit(); 184 | assert_eq!(fused.next().await, None); 185 | assert_eq!(fused.next().await, None); 186 | drop(fused); 187 | tx.until_empty().await; 188 | assert_eq!(tx.len(), 0); 189 | } 190 | { 191 | let mut tx = StopBroadcaster::new(); 192 | let rx = tx.listener(); 193 | let mut fused = rx.fuse_with(futures::stream::repeat(0).take(1)); 194 | assert_eq!(fused.next().await, Some(0)); 195 | assert_eq!(fused.next().await, None); 196 | assert_eq!(fused.next().await, None); 197 | tx.emit(); 198 | assert_eq!(fused.next().await, None); 199 | assert_eq!(fused.next().await, None); 200 | drop(fused); 201 | tx.until_empty().await; 202 | assert_eq!(tx.len(), 0); 203 | } 204 | } 205 | } 206 | -------------------------------------------------------------------------------- /src/test_util.rs: -------------------------------------------------------------------------------- 1 | use futures::{future::BoxFuture, Future, FutureExt, Stream, StreamExt}; 2 | use tokio::time::error::Elapsed; 3 | 4 | use crate::{StopListener, TmResult}; 5 | 6 | pub async fn quickpoll>(f: F) -> Result { 7 | tokio::time::timeout(tokio::time::Duration::from_millis(10), f).await 8 | } 9 | 10 | pub async fn not_ready(f: impl Future) -> bool { 11 | quickpoll(f).await.is_err() 12 | } 13 | 14 | pub async fn ready(f: impl Future) -> bool { 15 | quickpoll(f).await.is_ok() 16 | } 17 | 18 | pub async fn ok_fut(f: impl Future) -> TmResult { 19 | f.await; 20 | Ok(()) 21 | } 22 | 23 | pub fn blocker(info: &str, stop_rx: StopListener) -> BoxFuture<'static, String> { 24 | let info = info.to_string(); 25 | async move { 26 | stop_rx.await; 27 | tracing::info!("stopped: {}", info); 28 | info 29 | } 30 | .boxed() 31 | } 32 | 33 | pub fn fused<'a>( 34 | info: &str, 35 | mut stream: impl Stream + Unpin + Send + 'a, 36 | ) -> BoxFuture<'a, String> { 37 | let info = info.to_string(); 38 | async move { 39 | stream.next().await; 40 | tracing::info!("stopped: {}", info); 41 | info 42 | } 43 | .boxed() 44 | } 45 | --------------------------------------------------------------------------------