├── .env ├── .github ├── dependabot.yml └── workflows │ └── ci.yaml ├── .gitignore ├── .pre-commit-config.yaml ├── CHANGELOG.md ├── Cargo.lock ├── Cargo.toml ├── LICENSE ├── Makefile ├── README.md ├── examples ├── hello.rs ├── ping.rs ├── pub.rs └── sub.rs ├── rust-toolchain.toml ├── src ├── bin │ ├── cli.rs │ └── server.rs ├── client │ ├── cli.rs │ ├── cmd.rs │ ├── mod.rs │ └── subscriber.rs ├── cmd │ ├── get.rs │ ├── mod.rs │ ├── ping.rs │ ├── publish.rs │ ├── set.rs │ ├── subscribe.rs │ ├── unknown.rs │ └── unsubscribe.rs ├── config.rs ├── connection │ ├── connect.rs │ ├── frame.rs │ ├── mod.rs │ └── parse.rs ├── consts.rs ├── error.rs ├── lib.rs ├── logger.rs ├── server │ ├── handler.rs │ ├── listener.rs │ ├── mod.rs │ └── shutdown.rs └── storage │ ├── db.rs │ ├── mod.rs │ ├── store.rs │ └── traits.rs └── tests ├── client.rs └── server.rs /.env: -------------------------------------------------------------------------------- 1 | LOG_LEVEL=INFO -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | - package-ecosystem: "cargo" 4 | # Look for `Cargo.toml` and `Cargo.lock` in the root directory 5 | directory: "/" 6 | # Check for updates every Monday 7 | schedule: 8 | interval: "weekly" 9 | open-pull-requests-limit: 10 10 | - package-ecosystem: "github-actions" 11 | directory: "/" 12 | # Check for updates every Monday 13 | schedule: 14 | interval: "weekly" 15 | open-pull-requests-limit: 10 16 | -------------------------------------------------------------------------------- /.github/workflows/ci.yaml: -------------------------------------------------------------------------------- 1 | name: CI # Continuous Integration 2 | 3 | on: 4 | workflow_dispatch: 5 | push: 6 | paths-ignore: 7 | - '**.md' 8 | pull_request: 9 | paths-ignore: 10 | - '**.md' 11 | 12 | env: 13 | RUST_TOOLCHAIN: stable 14 | TOOLCHAIN_PROFILE: minimal 15 | 16 | jobs: 17 | lints: 18 | name: Run cargo fmt and cargo clippy 19 | runs-on: ubuntu-latest 20 | steps: 21 | - name: Checkout sources 22 | uses: actions/checkout@v3 23 | - name: Install toolchain 24 | uses: actions-rs/toolchain@v1 25 | with: 26 | profile: ${{ env.TOOLCHAIN_PROFILE }} 27 | toolchain: ${{ env.RUST_TOOLCHAIN }} 28 | override: true 29 | components: rustfmt, clippy 30 | - name: Cache 31 | uses: Swatinem/rust-cache@v2 32 | - name: Run cargo fmt 33 | uses: actions-rs/cargo@v1 34 | with: 35 | command: fmt 36 | args: --all -- --check 37 | - name: Run cargo clippy 38 | uses: actions-rs/cargo@v1 39 | with: 40 | command: clippy 41 | args: -- -D warnings 42 | test: 43 | name: Run cargo test 44 | runs-on: ubuntu-latest 45 | steps: 46 | - name: Checkout sources 47 | uses: actions/checkout@v3 48 | - name: Install toolchain 49 | uses: actions-rs/toolchain@v1 50 | with: 51 | profile: ${{ env.TOOLCHAIN_PROFILE }} 52 | toolchain: ${{ env.RUST_TOOLCHAIN }} 53 | override: true 54 | - name: Cache 55 | uses: Swatinem/rust-cache@v2 56 | - name: Run cargo test 57 | uses: actions-rs/cargo@v1 58 | env: 59 | RUST_TEST_THREADS: 8 60 | with: 61 | command: test 62 | args: --all-features 63 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | ### Rust template 2 | # Generated by Cargo 3 | # will have compiled files and executables 4 | debug/ 5 | target/ 6 | 7 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 8 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 9 | #Cargo.lock 10 | 11 | 12 | # These are backup files generated by rustfmt 13 | **/*.rs.bk 14 | 15 | ### Example user template template 16 | ### Example user template 17 | 18 | # IntelliJ project files 19 | .idea 20 | *.iml 21 | out 22 | gen 23 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pre-commit/pre-commit-hooks 3 | rev: v4.0.1 4 | hooks: 5 | - id: check-merge-conflict 6 | - id: check-toml 7 | - id: check-yaml 8 | - id: end-of-file-fixer 9 | - id: trailing-whitespace 10 | args: [ --markdown-linebreak-ext=md ] 11 | 12 | - repo: local 13 | hooks: 14 | - id: make-fmt 15 | name: make fmt 16 | entry: make fmt 17 | language: system 18 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | All notable changes to this project will be documented in this file. 3 | 4 | The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), 5 | and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). 6 | 7 | ## [Unreleased] 8 | -------------------------------------------------------------------------------- /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 = "async-stream" 7 | version = "0.3.3" 8 | source = "registry+https://github.com/rust-lang/crates.io-index" 9 | checksum = "dad5c83079eae9969be7fadefe640a1c566901f05ff91ab221de4b6f68d9507e" 10 | dependencies = [ 11 | "async-stream-impl", 12 | "futures-core", 13 | ] 14 | 15 | [[package]] 16 | name = "async-stream-impl" 17 | version = "0.3.3" 18 | source = "registry+https://github.com/rust-lang/crates.io-index" 19 | checksum = "10f203db73a71dfa2fb6dd22763990fa26f3d2625a6da2da900d23b87d26be27" 20 | dependencies = [ 21 | "proc-macro2", 22 | "quote", 23 | "syn", 24 | ] 25 | 26 | [[package]] 27 | name = "atoi" 28 | version = "2.0.0" 29 | source = "registry+https://github.com/rust-lang/crates.io-index" 30 | checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" 31 | dependencies = [ 32 | "num-traits", 33 | ] 34 | 35 | [[package]] 36 | name = "atty" 37 | version = "0.2.14" 38 | source = "registry+https://github.com/rust-lang/crates.io-index" 39 | checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" 40 | dependencies = [ 41 | "hermit-abi", 42 | "libc", 43 | "winapi", 44 | ] 45 | 46 | [[package]] 47 | name = "autocfg" 48 | version = "1.1.0" 49 | source = "registry+https://github.com/rust-lang/crates.io-index" 50 | checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" 51 | 52 | [[package]] 53 | name = "bitflags" 54 | version = "1.3.2" 55 | source = "registry+https://github.com/rust-lang/crates.io-index" 56 | checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" 57 | 58 | [[package]] 59 | name = "bytes" 60 | version = "1.3.0" 61 | source = "registry+https://github.com/rust-lang/crates.io-index" 62 | checksum = "dfb24e866b15a1af2a1b663f10c6b6b8f397a84aadb828f12e5b289ec23a3a3c" 63 | 64 | [[package]] 65 | name = "cfg-if" 66 | version = "1.0.0" 67 | source = "registry+https://github.com/rust-lang/crates.io-index" 68 | checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" 69 | 70 | [[package]] 71 | name = "clap" 72 | version = "3.2.23" 73 | source = "registry+https://github.com/rust-lang/crates.io-index" 74 | checksum = "71655c45cb9845d3270c9d6df84ebe72b4dad3c2ba3f7023ad47c144e4e473a5" 75 | dependencies = [ 76 | "atty", 77 | "bitflags", 78 | "clap_derive", 79 | "clap_lex", 80 | "indexmap", 81 | "once_cell", 82 | "strsim", 83 | "termcolor", 84 | "textwrap", 85 | ] 86 | 87 | [[package]] 88 | name = "clap_derive" 89 | version = "3.2.18" 90 | source = "registry+https://github.com/rust-lang/crates.io-index" 91 | checksum = "ea0c8bce528c4be4da13ea6fead8965e95b6073585a2f05204bd8f4119f82a65" 92 | dependencies = [ 93 | "heck", 94 | "proc-macro-error", 95 | "proc-macro2", 96 | "quote", 97 | "syn", 98 | ] 99 | 100 | [[package]] 101 | name = "clap_lex" 102 | version = "0.2.4" 103 | source = "registry+https://github.com/rust-lang/crates.io-index" 104 | checksum = "2850f2f5a82cbf437dd5af4d49848fbdfc27c157c3d010345776f952765261c5" 105 | dependencies = [ 106 | "os_str_bytes", 107 | ] 108 | 109 | [[package]] 110 | name = "dotenv" 111 | version = "0.15.0" 112 | source = "registry+https://github.com/rust-lang/crates.io-index" 113 | checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f" 114 | 115 | [[package]] 116 | name = "futures-core" 117 | version = "0.3.25" 118 | source = "registry+https://github.com/rust-lang/crates.io-index" 119 | checksum = "04909a7a7e4633ae6c4a9ab280aeb86da1236243a77b694a49eacd659a4bd3ac" 120 | 121 | [[package]] 122 | name = "getrandom" 123 | version = "0.2.8" 124 | source = "registry+https://github.com/rust-lang/crates.io-index" 125 | checksum = "c05aeb6a22b8f62540c194aac980f2115af067bfe15a0734d7277a768d396b31" 126 | dependencies = [ 127 | "cfg-if", 128 | "libc", 129 | "wasi", 130 | ] 131 | 132 | [[package]] 133 | name = "hashbrown" 134 | version = "0.12.3" 135 | source = "registry+https://github.com/rust-lang/crates.io-index" 136 | checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" 137 | 138 | [[package]] 139 | name = "heck" 140 | version = "0.4.0" 141 | source = "registry+https://github.com/rust-lang/crates.io-index" 142 | checksum = "2540771e65fc8cb83cd6e8a237f70c319bd5c29f78ed1084ba5d50eeac86f7f9" 143 | 144 | [[package]] 145 | name = "hermit-abi" 146 | version = "0.1.19" 147 | source = "registry+https://github.com/rust-lang/crates.io-index" 148 | checksum = "62b467343b94ba476dcb2500d242dadbb39557df889310ac77c5d99100aaac33" 149 | dependencies = [ 150 | "libc", 151 | ] 152 | 153 | [[package]] 154 | name = "indexmap" 155 | version = "1.9.2" 156 | source = "registry+https://github.com/rust-lang/crates.io-index" 157 | checksum = "1885e79c1fc4b10f0e172c475f458b7f7b93061064d98c3293e98c5ba0c8b399" 158 | dependencies = [ 159 | "autocfg", 160 | "hashbrown", 161 | ] 162 | 163 | [[package]] 164 | name = "libc" 165 | version = "0.2.137" 166 | source = "registry+https://github.com/rust-lang/crates.io-index" 167 | checksum = "fc7fcc620a3bff7cdd7a365be3376c97191aeaccc2a603e600951e452615bf89" 168 | 169 | [[package]] 170 | name = "lock_api" 171 | version = "0.4.9" 172 | source = "registry+https://github.com/rust-lang/crates.io-index" 173 | checksum = "435011366fe56583b16cf956f9df0095b405b82d76425bc8981c0e22e60ec4df" 174 | dependencies = [ 175 | "autocfg", 176 | "scopeguard", 177 | ] 178 | 179 | [[package]] 180 | name = "log" 181 | version = "0.4.17" 182 | source = "registry+https://github.com/rust-lang/crates.io-index" 183 | checksum = "abb12e687cfb44aa40f41fc3978ef76448f9b6038cad6aef4259d3c095a2382e" 184 | dependencies = [ 185 | "cfg-if", 186 | ] 187 | 188 | [[package]] 189 | name = "memchr" 190 | version = "2.5.0" 191 | source = "registry+https://github.com/rust-lang/crates.io-index" 192 | checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" 193 | 194 | [[package]] 195 | name = "mini-redis" 196 | version = "0.1.0" 197 | dependencies = [ 198 | "async-stream", 199 | "atoi", 200 | "bytes", 201 | "clap", 202 | "dotenv", 203 | "log", 204 | "rand", 205 | "thiserror", 206 | "tokio", 207 | "tokio-stream", 208 | ] 209 | 210 | [[package]] 211 | name = "mio" 212 | version = "0.8.5" 213 | source = "registry+https://github.com/rust-lang/crates.io-index" 214 | checksum = "e5d732bc30207a6423068df043e3d02e0735b155ad7ce1a6f76fe2baa5b158de" 215 | dependencies = [ 216 | "libc", 217 | "log", 218 | "wasi", 219 | "windows-sys", 220 | ] 221 | 222 | [[package]] 223 | name = "num-traits" 224 | version = "0.2.15" 225 | source = "registry+https://github.com/rust-lang/crates.io-index" 226 | checksum = "578ede34cf02f8924ab9447f50c28075b4d3e5b269972345e7e0372b38c6cdcd" 227 | dependencies = [ 228 | "autocfg", 229 | ] 230 | 231 | [[package]] 232 | name = "num_cpus" 233 | version = "1.14.0" 234 | source = "registry+https://github.com/rust-lang/crates.io-index" 235 | checksum = "f6058e64324c71e02bc2b150e4f3bc8286db6c83092132ffa3f6b1eab0f9def5" 236 | dependencies = [ 237 | "hermit-abi", 238 | "libc", 239 | ] 240 | 241 | [[package]] 242 | name = "once_cell" 243 | version = "1.16.0" 244 | source = "registry+https://github.com/rust-lang/crates.io-index" 245 | checksum = "86f0b0d4bf799edbc74508c1e8bf170ff5f41238e5f8225603ca7caaae2b7860" 246 | 247 | [[package]] 248 | name = "os_str_bytes" 249 | version = "6.4.1" 250 | source = "registry+https://github.com/rust-lang/crates.io-index" 251 | checksum = "9b7820b9daea5457c9f21c69448905d723fbd21136ccf521748f23fd49e723ee" 252 | 253 | [[package]] 254 | name = "parking_lot" 255 | version = "0.12.1" 256 | source = "registry+https://github.com/rust-lang/crates.io-index" 257 | checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" 258 | dependencies = [ 259 | "lock_api", 260 | "parking_lot_core", 261 | ] 262 | 263 | [[package]] 264 | name = "parking_lot_core" 265 | version = "0.9.5" 266 | source = "registry+https://github.com/rust-lang/crates.io-index" 267 | checksum = "7ff9f3fef3968a3ec5945535ed654cb38ff72d7495a25619e2247fb15a2ed9ba" 268 | dependencies = [ 269 | "cfg-if", 270 | "libc", 271 | "redox_syscall", 272 | "smallvec", 273 | "windows-sys", 274 | ] 275 | 276 | [[package]] 277 | name = "pin-project-lite" 278 | version = "0.2.9" 279 | source = "registry+https://github.com/rust-lang/crates.io-index" 280 | checksum = "e0a7ae3ac2f1173085d398531c705756c94a4c56843785df85a60c1a0afac116" 281 | 282 | [[package]] 283 | name = "ppv-lite86" 284 | version = "0.2.17" 285 | source = "registry+https://github.com/rust-lang/crates.io-index" 286 | checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" 287 | 288 | [[package]] 289 | name = "proc-macro-error" 290 | version = "1.0.4" 291 | source = "registry+https://github.com/rust-lang/crates.io-index" 292 | checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" 293 | dependencies = [ 294 | "proc-macro-error-attr", 295 | "proc-macro2", 296 | "quote", 297 | "syn", 298 | "version_check", 299 | ] 300 | 301 | [[package]] 302 | name = "proc-macro-error-attr" 303 | version = "1.0.4" 304 | source = "registry+https://github.com/rust-lang/crates.io-index" 305 | checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" 306 | dependencies = [ 307 | "proc-macro2", 308 | "quote", 309 | "version_check", 310 | ] 311 | 312 | [[package]] 313 | name = "proc-macro2" 314 | version = "1.0.47" 315 | source = "registry+https://github.com/rust-lang/crates.io-index" 316 | checksum = "5ea3d908b0e36316caf9e9e2c4625cdde190a7e6f440d794667ed17a1855e725" 317 | dependencies = [ 318 | "unicode-ident", 319 | ] 320 | 321 | [[package]] 322 | name = "quote" 323 | version = "1.0.21" 324 | source = "registry+https://github.com/rust-lang/crates.io-index" 325 | checksum = "bbe448f377a7d6961e30f5955f9b8d106c3f5e449d493ee1b125c1d43c2b5179" 326 | dependencies = [ 327 | "proc-macro2", 328 | ] 329 | 330 | [[package]] 331 | name = "rand" 332 | version = "0.8.5" 333 | source = "registry+https://github.com/rust-lang/crates.io-index" 334 | checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" 335 | dependencies = [ 336 | "libc", 337 | "rand_chacha", 338 | "rand_core", 339 | ] 340 | 341 | [[package]] 342 | name = "rand_chacha" 343 | version = "0.3.1" 344 | source = "registry+https://github.com/rust-lang/crates.io-index" 345 | checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" 346 | dependencies = [ 347 | "ppv-lite86", 348 | "rand_core", 349 | ] 350 | 351 | [[package]] 352 | name = "rand_core" 353 | version = "0.6.4" 354 | source = "registry+https://github.com/rust-lang/crates.io-index" 355 | checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" 356 | dependencies = [ 357 | "getrandom", 358 | ] 359 | 360 | [[package]] 361 | name = "redox_syscall" 362 | version = "0.2.16" 363 | source = "registry+https://github.com/rust-lang/crates.io-index" 364 | checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" 365 | dependencies = [ 366 | "bitflags", 367 | ] 368 | 369 | [[package]] 370 | name = "scopeguard" 371 | version = "1.1.0" 372 | source = "registry+https://github.com/rust-lang/crates.io-index" 373 | checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" 374 | 375 | [[package]] 376 | name = "signal-hook-registry" 377 | version = "1.4.0" 378 | source = "registry+https://github.com/rust-lang/crates.io-index" 379 | checksum = "e51e73328dc4ac0c7ccbda3a494dfa03df1de2f46018127f60c693f2648455b0" 380 | dependencies = [ 381 | "libc", 382 | ] 383 | 384 | [[package]] 385 | name = "smallvec" 386 | version = "1.10.0" 387 | source = "registry+https://github.com/rust-lang/crates.io-index" 388 | checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0" 389 | 390 | [[package]] 391 | name = "socket2" 392 | version = "0.4.7" 393 | source = "registry+https://github.com/rust-lang/crates.io-index" 394 | checksum = "02e2d2db9033d13a1567121ddd7a095ee144db4e1ca1b1bda3419bc0da294ebd" 395 | dependencies = [ 396 | "libc", 397 | "winapi", 398 | ] 399 | 400 | [[package]] 401 | name = "strsim" 402 | version = "0.10.0" 403 | source = "registry+https://github.com/rust-lang/crates.io-index" 404 | checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" 405 | 406 | [[package]] 407 | name = "syn" 408 | version = "1.0.104" 409 | source = "registry+https://github.com/rust-lang/crates.io-index" 410 | checksum = "4ae548ec36cf198c0ef7710d3c230987c2d6d7bd98ad6edc0274462724c585ce" 411 | dependencies = [ 412 | "proc-macro2", 413 | "quote", 414 | "unicode-ident", 415 | ] 416 | 417 | [[package]] 418 | name = "termcolor" 419 | version = "1.1.3" 420 | source = "registry+https://github.com/rust-lang/crates.io-index" 421 | checksum = "bab24d30b911b2376f3a13cc2cd443142f0c81dda04c118693e35b3835757755" 422 | dependencies = [ 423 | "winapi-util", 424 | ] 425 | 426 | [[package]] 427 | name = "textwrap" 428 | version = "0.16.0" 429 | source = "registry+https://github.com/rust-lang/crates.io-index" 430 | checksum = "222a222a5bfe1bba4a77b45ec488a741b3cb8872e5e499451fd7d0129c9c7c3d" 431 | 432 | [[package]] 433 | name = "thiserror" 434 | version = "1.0.38" 435 | source = "registry+https://github.com/rust-lang/crates.io-index" 436 | checksum = "6a9cd18aa97d5c45c6603caea1da6628790b37f7a34b6ca89522331c5180fed0" 437 | dependencies = [ 438 | "thiserror-impl", 439 | ] 440 | 441 | [[package]] 442 | name = "thiserror-impl" 443 | version = "1.0.38" 444 | source = "registry+https://github.com/rust-lang/crates.io-index" 445 | checksum = "1fb327af4685e4d03fa8cbcf1716380da910eeb2bb8be417e7f9fd3fb164f36f" 446 | dependencies = [ 447 | "proc-macro2", 448 | "quote", 449 | "syn", 450 | ] 451 | 452 | [[package]] 453 | name = "tokio" 454 | version = "1.23.0" 455 | source = "registry+https://github.com/rust-lang/crates.io-index" 456 | checksum = "eab6d665857cc6ca78d6e80303a02cea7a7851e85dfbd77cbdc09bd129f1ef46" 457 | dependencies = [ 458 | "autocfg", 459 | "bytes", 460 | "libc", 461 | "memchr", 462 | "mio", 463 | "num_cpus", 464 | "parking_lot", 465 | "pin-project-lite", 466 | "signal-hook-registry", 467 | "socket2", 468 | "tokio-macros", 469 | "windows-sys", 470 | ] 471 | 472 | [[package]] 473 | name = "tokio-macros" 474 | version = "1.8.0" 475 | source = "registry+https://github.com/rust-lang/crates.io-index" 476 | checksum = "9724f9a975fb987ef7a3cd9be0350edcbe130698af5b8f7a631e23d42d052484" 477 | dependencies = [ 478 | "proc-macro2", 479 | "quote", 480 | "syn", 481 | ] 482 | 483 | [[package]] 484 | name = "tokio-stream" 485 | version = "0.1.11" 486 | source = "registry+https://github.com/rust-lang/crates.io-index" 487 | checksum = "d660770404473ccd7bc9f8b28494a811bc18542b915c0855c51e8f419d5223ce" 488 | dependencies = [ 489 | "futures-core", 490 | "pin-project-lite", 491 | "tokio", 492 | ] 493 | 494 | [[package]] 495 | name = "unicode-ident" 496 | version = "1.0.5" 497 | source = "registry+https://github.com/rust-lang/crates.io-index" 498 | checksum = "6ceab39d59e4c9499d4e5a8ee0e2735b891bb7308ac83dfb4e80cad195c9f6f3" 499 | 500 | [[package]] 501 | name = "version_check" 502 | version = "0.9.4" 503 | source = "registry+https://github.com/rust-lang/crates.io-index" 504 | checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" 505 | 506 | [[package]] 507 | name = "wasi" 508 | version = "0.11.0+wasi-snapshot-preview1" 509 | source = "registry+https://github.com/rust-lang/crates.io-index" 510 | checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" 511 | 512 | [[package]] 513 | name = "winapi" 514 | version = "0.3.9" 515 | source = "registry+https://github.com/rust-lang/crates.io-index" 516 | checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" 517 | dependencies = [ 518 | "winapi-i686-pc-windows-gnu", 519 | "winapi-x86_64-pc-windows-gnu", 520 | ] 521 | 522 | [[package]] 523 | name = "winapi-i686-pc-windows-gnu" 524 | version = "0.4.0" 525 | source = "registry+https://github.com/rust-lang/crates.io-index" 526 | checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" 527 | 528 | [[package]] 529 | name = "winapi-util" 530 | version = "0.1.5" 531 | source = "registry+https://github.com/rust-lang/crates.io-index" 532 | checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" 533 | dependencies = [ 534 | "winapi", 535 | ] 536 | 537 | [[package]] 538 | name = "winapi-x86_64-pc-windows-gnu" 539 | version = "0.4.0" 540 | source = "registry+https://github.com/rust-lang/crates.io-index" 541 | checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" 542 | 543 | [[package]] 544 | name = "windows-sys" 545 | version = "0.42.0" 546 | source = "registry+https://github.com/rust-lang/crates.io-index" 547 | checksum = "5a3e1820f08b8513f676f7ab6c1f99ff312fb97b553d30ff4dd86f9f15728aa7" 548 | dependencies = [ 549 | "windows_aarch64_gnullvm", 550 | "windows_aarch64_msvc", 551 | "windows_i686_gnu", 552 | "windows_i686_msvc", 553 | "windows_x86_64_gnu", 554 | "windows_x86_64_gnullvm", 555 | "windows_x86_64_msvc", 556 | ] 557 | 558 | [[package]] 559 | name = "windows_aarch64_gnullvm" 560 | version = "0.42.0" 561 | source = "registry+https://github.com/rust-lang/crates.io-index" 562 | checksum = "41d2aa71f6f0cbe00ae5167d90ef3cfe66527d6f613ca78ac8024c3ccab9a19e" 563 | 564 | [[package]] 565 | name = "windows_aarch64_msvc" 566 | version = "0.42.0" 567 | source = "registry+https://github.com/rust-lang/crates.io-index" 568 | checksum = "dd0f252f5a35cac83d6311b2e795981f5ee6e67eb1f9a7f64eb4500fbc4dcdb4" 569 | 570 | [[package]] 571 | name = "windows_i686_gnu" 572 | version = "0.42.0" 573 | source = "registry+https://github.com/rust-lang/crates.io-index" 574 | checksum = "fbeae19f6716841636c28d695375df17562ca208b2b7d0dc47635a50ae6c5de7" 575 | 576 | [[package]] 577 | name = "windows_i686_msvc" 578 | version = "0.42.0" 579 | source = "registry+https://github.com/rust-lang/crates.io-index" 580 | checksum = "84c12f65daa39dd2babe6e442988fc329d6243fdce47d7d2d155b8d874862246" 581 | 582 | [[package]] 583 | name = "windows_x86_64_gnu" 584 | version = "0.42.0" 585 | source = "registry+https://github.com/rust-lang/crates.io-index" 586 | checksum = "bf7b1b21b5362cbc318f686150e5bcea75ecedc74dd157d874d754a2ca44b0ed" 587 | 588 | [[package]] 589 | name = "windows_x86_64_gnullvm" 590 | version = "0.42.0" 591 | source = "registry+https://github.com/rust-lang/crates.io-index" 592 | checksum = "09d525d2ba30eeb3297665bd434a54297e4170c7f1a44cad4ef58095b4cd2028" 593 | 594 | [[package]] 595 | name = "windows_x86_64_msvc" 596 | version = "0.42.0" 597 | source = "registry+https://github.com/rust-lang/crates.io-index" 598 | checksum = "f40009d85759725a34da6d89a94e63d7bdc50a862acf0dbc7c8e488f1edcb6f5" 599 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "mini-redis" 3 | version = "0.1.0" 4 | edition = "2018" 5 | description = "A mini-redis learn from tokio." 6 | repository = "https://github.com/JasonkayZK/mini-redis" 7 | license-file = "LICENSE" 8 | 9 | [[bin]] 10 | name = "mini-redis-cli" 11 | path = "src/bin/cli.rs" 12 | 13 | [[bin]] 14 | name = "mini-redis-server" 15 | path = "src/bin/server.rs" 16 | 17 | [dependencies] 18 | async-stream = "0.3.0" 19 | atoi = "2.0.0" 20 | bytes = "1" 21 | rand = "0.8.5" 22 | clap = { version = "3.2.23", features = ["derive"] } 23 | tokio = { version = "1", features = ["full"] } 24 | tokio-stream = "0.1" 25 | thiserror = "1.0.38" 26 | log = "0.4" 27 | dotenv = "0.15" 28 | 29 | [dev-dependencies] 30 | # Enable test-utilities in dev mode only. This is mostly for tests. 31 | tokio = { version = "1", features = ["test-util"] } 32 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 JasonkayZK <271226192@qq.com> 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: fmt clippy clean build pack all ci 2 | 3 | all: clean fmt clippy pack 4 | 5 | ci: fmt clippy 6 | 7 | fmt: 8 | cargo fmt --all -- 9 | 10 | clippy: 11 | cargo clippy -- -D warnings 12 | 13 | clean: 14 | rm -rf ./target 15 | 16 | build: 17 | cargo build 18 | 19 | pack: 20 | cargo build --release 21 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # mini-redis 2 | 3 | [![CI](https://github.com/JasonkayZK/mini-redis/workflows/CI/badge.svg)](https://github.com/JasonkayZK/mini-redis/actions) 4 | 5 |
6 | 7 | ## **相关文章** 8 | 9 | 系列文章： 10 | 11 | - [《mini-redis项目-1-简介》](https://jasonkayzk.github.io/2022/12/05/mini-redis项目-1-简介/) 12 | - [《mini-redis项目-2-存储层》](https://jasonkayzk.github.io/2022/12/05/mini-redis项目-2-存储层/) 13 | - [《mini-redis项目-3-连接层》](https://jasonkayzk.github.io/2022/12/05/mini-redis项目-3-连接层/) 14 | - [《mini-redis项目-4-服务端》](https://jasonkayzk.github.io/2022/12/06/mini-redis项目-4-服务端/) 15 | - [《mini-redis项目-5-客户端》](https://jasonkayzk.github.io/2022/12/07/mini-redis项目-5-客户端/) 16 | - [《mini-redis项目-6-测试与示例》](https://jasonkayzk.github.io/2022/12/07/mini-redis项目-6-测试与示例/) 17 | 18 | 19 | ## **前言** 20 | 21 | tokio 官方文档如下： 22 | 23 | - https://tokio.rs/tokio/tutorial 24 | 25 | 项目的目录结构如下： 26 | 27 | ```bash 28 | $ tree ./src/ 29 | . 30 | ├── bin 31 | │   ├── cli.rs 32 | │   └── server.rs 33 | ├── client 34 | │   ├── cli.rs 35 | │   ├── cmd.rs 36 | │   ├── mod.rs 37 | │   └── subscriber.rs 38 | ├── cmd 39 | │   ├── get.rs 40 | │   ├── mod.rs 41 | │   ├── ping.rs 42 | │   ├── publish.rs 43 | │   ├── set.rs 44 | │   ├── subscribe.rs 45 | │   ├── unknown.rs 46 | │   └── unsubscribe.rs 47 | ├── config.rs 48 | ├── connection 49 | │   ├── connect.rs 50 | │   ├── frame.rs 51 | │   ├── mod.rs 52 | │   └── parse.rs 53 | ├── consts.rs 54 | ├── error.rs 55 | ├── lib.rs 56 | ├── logger.rs 57 | ├── server 58 | │   ├── handler.rs 59 | │   ├── listener.rs 60 | │   ├── mod.rs 61 | │   └── shutdown.rs 62 | └── storage 63 | ├── db.rs 64 | ├── mod.rs 65 | ├── store.rs 66 | └── traits.rs 67 | ``` 68 | 69 | 其中： 70 | 71 | - `bin` 目录：server 和 cli 的命令行入口可执行文件； 72 | - `client` 目录：客户端具体实现逻辑； 73 | - `server` 目录：服务端具体实现逻辑； 74 | - `cmd` 目录：mini-redis 相关命令实现； 75 | - `connection` 目录：客户端、服务端异步连接实现； 76 | - `storage` 目录：kv、subscribe 存储实现（本例中直接使用 HashMap 实现，实际生产环境多用 LSM-Tree）； 77 | - `config.rs`：mini-redis 配置相关； 78 | - `consts.rs`：mini-redis 常量配置相关； 79 | - `error.rs`：mini-redis 错误定义； 80 | - `logger.rs`：mini-redis 日志配置； 81 | - `lib.rs`：mini-redis 库入口； 82 | 83 | 总体分为下面几个部分： 84 | 85 | - **存储实现；** 86 | - **连接实现；** 87 | - **具体命令实现** 88 | - **客户端、服务端实现；** 89 | 90 |
91 | 92 | ## **基本使用** 93 | 94 | 首先启动server： 95 | 96 | ```bash 97 | $ cargo run --bin mini-redis-server 98 | 99 | [ INFO]: mini_redis::server - mini-redis server started listen on: 0.0.0.0:6379 100 | [ INFO]: mini_redis::server::listener - server started, accepting inbound connections 101 | ``` 102 | 103 | 随后可以使用 client： 104 | 105 | ```bash 106 | $ cargo run --bin mini-redis-cli 107 | 108 | mini-redis-cli 0.1.0 109 | Issue Redis commands 110 | 111 | USAGE: 112 | mini-redis-cli [OPTIONS] 113 | 114 | OPTIONS: 115 | -h, --help Print help information 116 | --hostname [default: 127.0.0.1] 117 | --port [default: 6379] 118 | -V, --version Print version information 119 | 120 | SUBCOMMANDS: 121 | get Get the value of key 122 | help Print this message or the help of the given subcommand(s) 123 | ping 124 | publish Publisher to send a message to a specific channel 125 | set Set key to hold the string value 126 | subscribe Subscribe a client to a specific channel or channels 127 | ``` 128 | 129 |
130 | 131 | ping命令测试： 132 | 133 | ```bash 134 | $ cargo run --bin mini-redis-cli ping 135 | "PONG" 136 | 137 | $ cargo run --bin mini-redis-cli ping abc 138 | "abc" 139 | ``` 140 | 141 |
142 | 143 | get/set 测试： 144 | 145 | ```bash 146 | $ cargo run --bin mini-redis-cli get foo 147 | (nil) 148 | 149 | $ cargo run --bin mini-redis-cli set foo 123 150 | OK 151 | 152 | $ cargo run --bin mini-redis-cli get foo 153 | "123" 154 | ``` 155 | 156 | 过期键测试，设置 5s 过期： 157 | 158 | ```bash 159 | $ cargo run --bin mini-redis-cli set foo 123 5000 160 | ``` 161 | 162 | 获取： 163 | 164 | ```bash 165 | $ cargo run --bin mini-redis-cli get foo 166 | "123" 167 | 168 | $ cargo run --bin mini-redis-cli get foo 169 | (nil) 170 | ``` 171 | 172 | 5s后，获取不到 key 值了！ 173 | 174 |
175 | 176 | pub/sub 测试； 177 | 178 | 启动三个 subscribe，订阅同一个 channel，ch1： 179 | 180 | ```bash 181 | $ cargo run --bin mini-redis-cli subscribe ch1 182 | 183 | $ cargo run --bin mini-redis-cli subscribe ch1 184 | 185 | $ cargo run --bin mini-redis-cli subscribe ch1 186 | ``` 187 | 188 | 向 ch1 发布消息： 189 | 190 | ```bash 191 | $ cargo run --bin mini-redis-cli publish ch1 a-message 192 | Publish OK 193 | ``` 194 | 195 | 其他订阅者均收到消息： 196 | 197 | ``` 198 | got message from the channel: ch1; message = b"a-message" 199 | ``` 200 | 201 |
202 | 203 | 错误命令测试： 204 | 205 | ```bash 206 | $ cargo run --bin mini-redis-cli ping get foo 207 | 208 | error: Found argument 'foo' which wasn't expected, or isn't valid in this context 209 | ``` 210 | -------------------------------------------------------------------------------- /examples/hello.rs: -------------------------------------------------------------------------------- 1 | //! Hello world client. 2 | //! 3 | //! A simple client that connects to a mini-redis server, sets key "hello" with value "world", 4 | //! and gets it from the server after 5 | //! 6 | //! You can test this out by running: 7 | //! 8 | //! cargo run --bin mini-redis-server 9 | //! 10 | //! And then in another terminal run: 11 | //! 12 | //! cargo run --example hello_world 13 | 14 | use mini_redis::client; 15 | use mini_redis::error::MiniRedisClientError; 16 | 17 | #[tokio::main] 18 | pub async fn main() -> Result<(), MiniRedisClientError> { 19 | // Open a connection to the mini-redis address. 20 | let mut client = client::connect("127.0.0.1:6379").await?; 21 | 22 | // Set the key "hello" with value "world" 23 | let result = client.set("hello", "world".into()).await?; 24 | println!("set value to the server success， result: {:?}", result); 25 | 26 | // Get key "hello" 27 | let result = client.get("hello").await?; 28 | println!("got value from the server success， result: {:?}", result); 29 | 30 | Ok(()) 31 | } 32 | -------------------------------------------------------------------------------- /examples/ping.rs: -------------------------------------------------------------------------------- 1 | //! Ping client. 2 | //! 3 | //! A simple client that connects to a mini-redis server, and say `ping` 4 | //! 5 | //! You can test this out by running: 6 | //! 7 | //! cargo run --bin mini-redis-server 8 | //! 9 | //! And then in another terminal run: 10 | //! 11 | //! cargo run --example ping 12 | 13 | use mini_redis::client; 14 | use mini_redis::error::MiniRedisConnectionError; 15 | 16 | #[tokio::main] 17 | pub async fn main() -> Result<(), MiniRedisConnectionError> { 18 | // Open a connection to the mini-redis address. 19 | let mut client = client::connect("127.0.0.1:6379").await?; 20 | 21 | let result = client.ping(None).await?; 22 | println!("empty ping response: {:?}", result); 23 | 24 | let result = client.ping(Some("hello".into())).await?; 25 | println!("bytes ping response: {:?}", result); 26 | 27 | Ok(()) 28 | } 29 | -------------------------------------------------------------------------------- /examples/pub.rs: -------------------------------------------------------------------------------- 1 | //! Publish to a redis channel example. 2 | //! 3 | //! A simple client that connects to a mini-redis server, and 4 | //! publishes a message on `foo` channel 5 | //! 6 | //! You can test this out by running: 7 | //! 8 | //! cargo run --bin mini-redis-server 9 | //! 10 | //! Then in another terminal run: 11 | //! 12 | //! cargo run --example sub 13 | //! 14 | //! And then in another terminal run: 15 | //! 16 | //! cargo run --example pub 17 | 18 | use mini_redis::client; 19 | use mini_redis::error::MiniRedisClientError; 20 | 21 | #[tokio::main] 22 | async fn main() -> Result<(), MiniRedisClientError> { 23 | // Open a connection to the mini-redis address. 24 | let mut client = client::connect("127.0.0.1:6379").await?; 25 | 26 | // publish message `bar` on channel foo 27 | let res = client.publish("foo", "bar".into()).await?; 28 | 29 | println!("pushed message success, res: {:?}", res); 30 | 31 | Ok(()) 32 | } 33 | -------------------------------------------------------------------------------- /examples/sub.rs: -------------------------------------------------------------------------------- 1 | //! Subscribe to a redis channel example. 2 | //! 3 | //! A simple client that connects to a mini-redis server, subscribes to "foo" and "bar" channels 4 | //! and awaits messages published on those channels 5 | //! 6 | //! You can test this out by running: 7 | //! 8 | //! cargo run --bin mini-redis-server 9 | //! 10 | //! Then in another terminal run: 11 | //! 12 | //! cargo run --example sub 13 | //! 14 | //! And then in another terminal run: 15 | //! 16 | //! cargo run --example pub 17 | 18 | use mini_redis::client; 19 | use mini_redis::error::MiniRedisClientError; 20 | 21 | #[tokio::main] 22 | pub async fn main() -> Result<(), MiniRedisClientError> { 23 | // Open a connection to the mini-redis address. 24 | let client = client::connect("127.0.0.1:6379").await?; 25 | 26 | // subscribe to channel foo 27 | let mut subscriber = client.subscribe(vec!["foo".into()]).await?; 28 | 29 | // await messages on channel foo 30 | if let Some(msg) = subscriber.next_message().await? { 31 | println!( 32 | "got message from the channel: {}; message = {:?}", 33 | msg.channel, msg.content 34 | ); 35 | } 36 | 37 | Ok(()) 38 | } 39 | -------------------------------------------------------------------------------- /rust-toolchain.toml: -------------------------------------------------------------------------------- 1 | [toolchain] 2 | channel = "stable" 3 | -------------------------------------------------------------------------------- /src/bin/cli.rs: -------------------------------------------------------------------------------- 1 | use clap::Parser; 2 | use dotenv::dotenv; 3 | use log::debug; 4 | 5 | use mini_redis::client::cmd::Command; 6 | use mini_redis::consts::DEFAULT_PORT; 7 | use mini_redis::error::{MiniRedisClientError, MiniRedisConnectionError}; 8 | use mini_redis::{client, logger}; 9 | 10 | #[derive(Parser, Debug)] 11 | #[clap( 12 | name = "mini-redis-cli", 13 | version, 14 | author, 15 | about = "Issue Redis commands" 16 | )] 17 | struct Cli { 18 | #[clap(subcommand)] 19 | command: Command, 20 | 21 | #[clap(name = "hostname", long, default_value = "127.0.0.1")] 22 | host: String, 23 | 24 | #[clap(long, default_value_t = DEFAULT_PORT)] 25 | port: u16, 26 | } 27 | 28 | /// Entry point for CLI tool. 29 | /// 30 | /// The `[tokio::main]` annotation signals that the Tokio runtime should be 31 | /// started when the function is called. The body of the function is executed 32 | /// within the newly spawned runtime. 33 | /// 34 | /// `flavor = "current_thread"` is used here to avoid spawning background 35 | /// threads. The CLI tool use case benefits more by being lighter instead of 36 | /// multi-threaded. 37 | #[tokio::main(flavor = "current_thread")] 38 | async fn main() -> Result<(), MiniRedisClientError> { 39 | dotenv().ok(); 40 | logger::init(); 41 | 42 | // Parse command line arguments 43 | let cli = Cli::parse(); 44 | debug!("get cli: {:?}", cli); 45 | 46 | // Get the remote address to connect to 47 | let addr = format!("{}:{}", cli.host, cli.port); 48 | 49 | // Establish a connection 50 | let mut client = client::connect(&addr).await?; 51 | 52 | // Process the requested command 53 | match cli.command { 54 | Command::Ping { msg } => { 55 | let value = client.ping(msg).await?; 56 | if let Ok(string) = std::str::from_utf8(&value) { 57 | println!("\"{}\"", string); 58 | } else { 59 | println!("{:?}", value); 60 | } 61 | } 62 | Command::Get { key } => { 63 | if let Some(value) = client.get(&key).await? { 64 | if let Ok(string) = std::str::from_utf8(&value) { 65 | println!("\"{}\"", string); 66 | } else { 67 | println!("{:?}", value); 68 | } 69 | } else { 70 | println!("(nil)"); 71 | } 72 | } 73 | Command::Set { 74 | key, 75 | value, 76 | expires: None, 77 | } => { 78 | client.set(&key, value).await?; 79 | println!("OK"); 80 | } 81 | Command::Set { 82 | key, 83 | value, 84 | expires: Some(expires), 85 | } => { 86 | client.set_expires(&key, value, expires).await?; 87 | println!("OK"); 88 | } 89 | Command::Publish { channel, message } => { 90 | client.publish(&channel, message).await?; 91 | println!("Publish OK"); 92 | } 93 | Command::Subscribe { channels } => { 94 | if channels.is_empty() { 95 | return Err(MiniRedisConnectionError::InvalidArgument( 96 | "channel(s) must be provided".into(), 97 | ) 98 | .into()); 99 | } 100 | let mut subscriber = client.subscribe(channels).await?; 101 | 102 | // await messages on channels 103 | while let Some(msg) = subscriber.next_message().await? { 104 | println!( 105 | "got message from the channel: {}; message = {:?}", 106 | msg.channel, msg.content 107 | ); 108 | } 109 | } 110 | } 111 | 112 | Ok(()) 113 | } 114 | -------------------------------------------------------------------------------- /src/bin/server.rs: -------------------------------------------------------------------------------- 1 | //! mini-redis server. 2 | //! 3 | //! This file is the entry point for the server implemented in the library. It 4 | //! performs command line parsing and passes the arguments on to 5 | //! `mini_redis::server`. 6 | //! 7 | //! The `clap` crate is used for parsing arguments. 8 | 9 | use clap::Parser; 10 | use dotenv::dotenv; 11 | use tokio::net::TcpListener; 12 | use tokio::signal; 13 | 14 | use mini_redis::consts::DEFAULT_PORT; 15 | use mini_redis::error::MiniRedisServerError; 16 | use mini_redis::{logger, server}; 17 | 18 | #[derive(Parser, Debug)] 19 | #[clap( 20 | name = "mini-redis-server", 21 | version, 22 | author, 23 | about = "A mini redis server" 24 | )] 25 | struct Cli { 26 | #[clap(long)] 27 | port: Option, 28 | } 29 | 30 | #[tokio::main] 31 | pub async fn main() -> Result<(), MiniRedisServerError> { 32 | let cli = init(); 33 | let port = cli.port.unwrap_or(DEFAULT_PORT); 34 | 35 | // Bind a TCP listener 36 | let listener = TcpListener::bind(&format!("0.0.0.0:{}", port)).await?; 37 | 38 | server::run(listener, signal::ctrl_c()).await; 39 | 40 | Ok(()) 41 | } 42 | 43 | fn init() -> Cli { 44 | dotenv().ok(); 45 | logger::init(); 46 | Cli::parse() 47 | } 48 | -------------------------------------------------------------------------------- /src/client/cli.rs: -------------------------------------------------------------------------------- 1 | //! Minimal Redis client implementation 2 | //! 3 | //! Provides an async connect and methods for issuing the supported commands. 4 | 5 | use crate::client::subscriber::Subscriber; 6 | use crate::cmd::get::Get; 7 | use bytes::Bytes; 8 | use log::{debug, error}; 9 | use std::time::Duration; 10 | 11 | use crate::cmd::ping::Ping; 12 | use crate::cmd::publish::Publish; 13 | use crate::cmd::set::Set; 14 | use crate::cmd::subscribe::Subscribe; 15 | use crate::connection::connect::Connection; 16 | use crate::connection::frame::Frame; 17 | use crate::error::MiniRedisConnectionError; 18 | 19 | /// Established connection with a Redis server. 20 | /// 21 | /// Backed by a single `TcpStream`, `Client` provides basic network client 22 | /// functionality (no pooling, retrying, ...). Connections are established using 23 | /// the [`connect`](fn@connect) function. 24 | /// 25 | /// Requests are issued using the various methods of `Client`. 26 | pub struct Client { 27 | /// The TCP connection decorated with the redis protocol encoder / decoder 28 | /// implemented using a buffered `TcpStream`. 29 | /// 30 | /// When `Listener` receives an inbound connection, the `TcpStream` is 31 | /// passed to `Connection::new`, which initializes the associated buffers. 32 | /// `Connection` allows the handler to operate at the "frame" level and keep 33 | /// the byte level protocol parsing details encapsulated in `Connection`. 34 | pub(crate) connection: Connection, 35 | } 36 | 37 | impl Client { 38 | /// Ping to the server. 39 | /// 40 | /// Returns PONG if no argument is provided, otherwise 41 | /// return a copy of the argument as a bulk. 42 | /// 43 | /// This command is often used to test if a connection 44 | /// is still alive, or to measure latency. 45 | /// 46 | /// # Examples 47 | /// 48 | /// Demonstrates basic usage. 49 | /// ```no_run 50 | /// 51 | /// #[tokio::main] 52 | /// async fn main() { 53 | /// let mut client = mini_redis::client::connect("localhost:6379").await.unwrap(); 54 | /// 55 | /// let pong = client.ping(None).await.unwrap(); 56 | /// assert_eq!(b"PONG", &pong[..]); 57 | /// } 58 | /// ``` 59 | pub async fn ping(&mut self, msg: Option) -> Result { 60 | let frame = Ping::new(msg).into_frame()?; 61 | debug!("request: {:?}", frame); 62 | 63 | self.connection.write_frame(&frame).await?; 64 | 65 | match self.read_response().await? { 66 | Frame::Simple(value) => Ok(value.into()), 67 | Frame::Bulk(value) => Ok(value), 68 | frame => Err(MiniRedisConnectionError::CommandExecute(frame.to_string())), 69 | } 70 | } 71 | 72 | /// Get the value of key. 73 | /// 74 | /// If the key does not exist the special value `None` is returned. 75 | /// 76 | /// # Examples 77 | /// 78 | /// Demonstrates basic usage. 79 | /// 80 | /// ```no_run 81 | /// #[tokio::main] 82 | /// async fn main() { 83 | /// let mut client = mini_redis::client::connect("localhost:6379").await.unwrap(); 84 | /// 85 | /// let val = client.get("foo").await.unwrap(); 86 | /// println!("Got = {:?}", val); 87 | /// } 88 | /// ``` 89 | pub async fn get(&mut self, key: &str) -> Result, MiniRedisConnectionError> { 90 | // Create a `Get` command for the `key` and convert it to a frame. 91 | let frame = Get::new(key).into_frame()?; 92 | 93 | debug!("get command request: {:?}", frame); 94 | 95 | // Write the frame to the socket. This writes the full frame to the 96 | // socket, waiting if necessary. 97 | self.connection.write_frame(&frame).await?; 98 | 99 | // Wait for the response from the server 100 | // 101 | // Both `Simple` and `Bulk` frames are accepted. `Null` represents the 102 | // key not being present and `None` is returned. 103 | match self.read_response().await? { 104 | Frame::Simple(value) => Ok(Some(value.into())), 105 | Frame::Bulk(value) => Ok(Some(value)), 106 | Frame::Null => Ok(None), 107 | frame => Err(MiniRedisConnectionError::CommandExecute(frame.to_string())), 108 | } 109 | } 110 | 111 | /// Set `key` to hold the given `value`. 112 | /// 113 | /// The `value` is associated with `key` until it is overwritten by the next 114 | /// call to `set` or it is removed. 115 | /// 116 | /// If key already holds a value, it is overwritten. Any previous time to 117 | /// live associated with the key is discarded on successful SET operation. 118 | /// 119 | /// # Examples 120 | /// 121 | /// Demonstrates basic usage. 122 | /// 123 | /// ```no_run 124 | /// #[tokio::main] 125 | /// async fn main() { 126 | /// let mut client = mini_redis::client::connect("localhost:6379").await.unwrap(); 127 | /// 128 | /// client.set("foo", "bar".into()).await.unwrap(); 129 | /// 130 | /// // Getting the value immediately works 131 | /// let val = client.get("foo").await.unwrap().unwrap(); 132 | /// assert_eq!(val, "bar"); 133 | /// } 134 | /// ``` 135 | pub async fn set(&mut self, key: &str, value: Bytes) -> Result<(), MiniRedisConnectionError> { 136 | // Create a `Set` command and pass it to `set_cmd`. A separate method is 137 | // used to set a value with an expiration. The common parts of both 138 | // functions are implemented by `set_cmd`. 139 | self.set_cmd(Set::new(key, value, None)).await 140 | } 141 | 142 | /// Set `key` to hold the given `value`. The value expires after `expiration` 143 | /// 144 | /// The `value` is associated with `key` until one of the following: 145 | /// - it expires. 146 | /// - it is overwritten by the next call to `set`. 147 | /// - it is removed. 148 | /// 149 | /// If key already holds a value, it is overwritten. Any previous time to 150 | /// live associated with the key is discarded on a successful SET operation. 151 | /// 152 | /// # Examples 153 | /// 154 | /// Demonstrates basic usage. This example is not **guaranteed** to always 155 | /// work as it relies on time based logic and assumes the client and server 156 | /// stay relatively synchronized in time. The real world tends to not be so 157 | /// favorable. 158 | /// 159 | /// ```no_run 160 | /// use tokio::time; 161 | /// use std::time::Duration; 162 | /// 163 | /// #[tokio::main] 164 | /// async fn main() { 165 | /// let ttl = Duration::from_millis(500); 166 | /// let mut client = mini_redis::client::connect("localhost:6379").await.unwrap(); 167 | /// 168 | /// client.set_expires("foo", "bar".into(), ttl).await.unwrap(); 169 | /// 170 | /// // Getting the value immediately works 171 | /// let val = client.get("foo").await.unwrap().unwrap(); 172 | /// assert_eq!(val, "bar"); 173 | /// 174 | /// // Wait for the TTL to expire 175 | /// time::sleep(ttl).await; 176 | /// 177 | /// let val = client.get("foo").await.unwrap(); 178 | /// assert!(val.is_some()); 179 | /// } 180 | /// ``` 181 | pub async fn set_expires( 182 | &mut self, 183 | key: &str, 184 | value: Bytes, 185 | expiration: Duration, 186 | ) -> Result<(), MiniRedisConnectionError> { 187 | // Create a `Set` command and pass it to `set_cmd`. A separate method is 188 | // used to set a value with an expiration. The common parts of both 189 | // functions are implemented by `set_cmd`. 190 | self.set_cmd(Set::new(key, value, Some(expiration))).await 191 | } 192 | 193 | /// The core `SET` logic, used by both `set` and `set_expires. 194 | async fn set_cmd(&mut self, cmd: Set) -> Result<(), MiniRedisConnectionError> { 195 | // Convert the `Set` command into a frame 196 | let frame = cmd.into_frame()?; 197 | 198 | debug!("set command request: {:?}", frame); 199 | 200 | // Write the frame to the socket. This writes the full frame to the 201 | // socket, waiting if necessary. 202 | self.connection.write_frame(&frame).await?; 203 | 204 | // Wait for the response from the server. On success, the server 205 | // responds simply with `OK`. Any other response indicates an error. 206 | match self.read_response().await? { 207 | Frame::Simple(response) if response == "OK" => Ok(()), 208 | frame => Err(MiniRedisConnectionError::CommandExecute(frame.to_string())), 209 | } 210 | } 211 | 212 | /// Posts `message` to the given `channel`. 213 | /// 214 | /// Returns the number of subscribers currently listening on the channel. 215 | /// There is no guarantee that these subscribers receive the message as they 216 | /// may disconnect at any time. 217 | /// 218 | /// # Examples 219 | /// 220 | /// Demonstrates basic usage. 221 | /// 222 | /// ```no_run 223 | /// #[tokio::main] 224 | /// async fn main() { 225 | /// let mut client = mini_redis::client::connect("localhost:6379").await.unwrap(); 226 | /// 227 | /// let val = client.publish("foo", "bar".into()).await.unwrap(); 228 | /// println!("Got = {:?}", val); 229 | /// } 230 | /// ``` 231 | pub async fn publish( 232 | &mut self, 233 | channel: &str, 234 | message: Bytes, 235 | ) -> Result { 236 | // Convert the `Publish` command into a frame 237 | let frame = Publish::new(channel, message).into_frame()?; 238 | 239 | debug!("publish command request: {:?}", frame); 240 | 241 | // Write the frame to the socket 242 | self.connection.write_frame(&frame).await?; 243 | 244 | // Read the response 245 | match self.read_response().await? { 246 | Frame::Integer(response) => Ok(response), 247 | frame => Err(MiniRedisConnectionError::CommandExecute(frame.to_string())), 248 | } 249 | } 250 | 251 | /// Subscribes the client to the specified channels. 252 | /// 253 | /// Once a client issues a subscribe command, it may no longer issue any 254 | /// non-pub/sub commands. The function consumes `self` and returns a `Subscriber`. 255 | /// 256 | /// The `Subscriber` value is used to receive messages as well as manage the 257 | /// list of channels the client is subscribed to. 258 | pub async fn subscribe( 259 | mut self, 260 | channels: Vec, 261 | ) -> Result { 262 | // Issue the subscribe command to the server and wait for confirmation. 263 | // The client will then have been transitioned into the "subscriber" 264 | // state and may only issue pub/sub commands from that point on. 265 | self.subscribe_cmd(&channels).await?; 266 | 267 | // Return the `Subscriber` type 268 | Ok(Subscriber { 269 | client: self, 270 | subscribed_channels: channels, 271 | }) 272 | } 273 | 274 | /// The core `SUBSCRIBE` logic, used by misc subscribe fns 275 | pub(crate) async fn subscribe_cmd( 276 | &mut self, 277 | channels: &[String], 278 | ) -> Result<(), MiniRedisConnectionError> { 279 | // Convert the `Subscribe` command into a frame 280 | let frame = Subscribe::new(channels).into_frame()?; 281 | 282 | debug!("subscribe command request: {:?}", frame); 283 | 284 | // Write the frame to the socket 285 | self.connection.write_frame(&frame).await?; 286 | 287 | // For each channel being subscribed to, the server responds with a 288 | // message confirming subscription to that channel. 289 | for channel in channels { 290 | // Read the response 291 | let response = self.read_response().await?; 292 | 293 | // Verify it is confirmation of subscription. 294 | match response { 295 | Frame::Array(ref frame) => match frame.as_slice() { 296 | // The server responds with an array frame in the form of: 297 | // 298 | // ``` 299 | // [ "subscribe", channel, num-subscribed ] 300 | // ``` 301 | // 302 | // where channel is the name of the channel and 303 | // num-subscribed is the number of channels that the client 304 | // is currently subscribed to. 305 | [subscribe, schannel, ..] 306 | if *subscribe == "subscribe" && *schannel == channel => 307 | { 308 | debug!("subscribe channel: {} success", channel); 309 | } 310 | _ => { 311 | error!("subscribe frame failed, response: {}", response); 312 | return Err(MiniRedisConnectionError::CommandExecute( 313 | response.to_string(), 314 | )); 315 | } 316 | }, 317 | frame => { 318 | error!( 319 | "subscribe frame failed, response frame type not match: {}", 320 | frame 321 | ); 322 | return Err(MiniRedisConnectionError::InvalidFrameType); 323 | } 324 | }; 325 | } 326 | 327 | Ok(()) 328 | } 329 | 330 | /// Reads a response frame from the socket. 331 | /// 332 | /// If an `Error` frame is received, it is converted to `Err`. 333 | pub(crate) async fn read_response(&mut self) -> Result { 334 | let response = self.connection.read_frame().await?; 335 | 336 | debug!("read response: {:?}", response); 337 | 338 | match response { 339 | // Error frames are converted to `Err` 340 | Some(Frame::Error(msg)) => Err(MiniRedisConnectionError::CommandExecute(msg)), 341 | Some(frame) => Ok(frame), 342 | None => { 343 | // Receiving `None` here indicates the server has closed the 344 | // connection without sending a frame. This is unexpected and is 345 | // represented as a "connection reset by peer" error. 346 | Err(MiniRedisConnectionError::Disconnect) 347 | } 348 | } 349 | } 350 | } 351 | -------------------------------------------------------------------------------- /src/client/cmd.rs: -------------------------------------------------------------------------------- 1 | use std::num::ParseIntError; 2 | use std::time::Duration; 3 | 4 | use bytes::Bytes; 5 | use clap::Subcommand; 6 | 7 | #[derive(Subcommand, Debug)] 8 | pub enum Command { 9 | Ping { 10 | /// Message to ping 11 | msg: Option, 12 | }, 13 | /// Get the value of key. 14 | Get { 15 | /// Name of key to get 16 | key: String, 17 | }, 18 | /// Set key to hold the string value. 19 | Set { 20 | /// Name of key to set 21 | key: String, 22 | 23 | /// Value to set. 24 | #[clap(parse(from_str = bytes_from_str))] 25 | value: Bytes, 26 | 27 | /// Expire the value after specified amount of time 28 | #[clap(parse(try_from_str = duration_from_ms_str))] 29 | expires: Option, 30 | }, 31 | /// Publisher to send a message to a specific channel. 32 | Publish { 33 | /// Name of channel 34 | channel: String, 35 | 36 | #[clap(parse(from_str = bytes_from_str))] 37 | /// Message to publish 38 | message: Bytes, 39 | }, 40 | /// Subscribe a client to a specific channel or channels. 41 | Subscribe { 42 | /// Specific channel or channels 43 | channels: Vec, 44 | }, 45 | } 46 | 47 | fn duration_from_ms_str(src: &str) -> Result { 48 | let ms = src.parse::()?; 49 | Ok(Duration::from_millis(ms)) 50 | } 51 | 52 | fn bytes_from_str(src: &str) -> Bytes { 53 | Bytes::from(src.to_string()) 54 | } 55 | -------------------------------------------------------------------------------- /src/client/mod.rs: -------------------------------------------------------------------------------- 1 | use tokio::net::{TcpStream, ToSocketAddrs}; 2 | 3 | use crate::client::cli::Client; 4 | use crate::connection::connect::Connection; 5 | use crate::error::MiniRedisConnectionError; 6 | 7 | pub mod cli; 8 | pub mod cmd; 9 | mod subscriber; 10 | 11 | /// Establish a connection with the Redis server located at `addr`. 12 | /// 13 | /// `addr` may be any type that can be asynchronously converted to a 14 | /// `SocketAddr`. This includes `SocketAddr` and strings. The `ToSocketAddrs` 15 | /// trait is the Tokio version and not the `std` version. 16 | /// 17 | /// # Examples 18 | /// 19 | /// ```no_run 20 | /// #[tokio::main] 21 | /// async fn main() { 22 | /// 23 | /// let client = match mini_redis::client::connect("localhost:6379").await { 24 | /// Ok(client) => client, 25 | /// Err(_) => panic!("failed to establish connection"), 26 | /// }; 27 | /// # drop(client); 28 | /// } 29 | /// ``` 30 | /// 31 | pub async fn connect(addr: T) -> Result { 32 | // The `addr` argument is passed directly to `TcpStream::connect`. This 33 | // performs any asynchronous DNS lookup and attempts to establish the TCP 34 | // connection. An error at either step returns an error, which is then 35 | // bubbled up to the caller of `mini_redis` connect. 36 | let socket = TcpStream::connect(addr).await?; 37 | 38 | // Initialize the connection state. This allocates read/write buffers to 39 | // perform redis protocol frame parsing. 40 | let connection = Connection::new(socket); 41 | 42 | Ok(Client { connection }) 43 | } 44 | -------------------------------------------------------------------------------- /src/client/subscriber.rs: -------------------------------------------------------------------------------- 1 | use async_stream::try_stream; 2 | use bytes::Bytes; 3 | use log::{debug, error}; 4 | use tokio_stream::Stream; 5 | 6 | use crate::client::cli::Client; 7 | use crate::cmd::unsubscribe::Unsubscribe; 8 | use crate::connection::frame::Frame; 9 | use crate::error::MiniRedisConnectionError; 10 | 11 | /// A client that has entered pub/sub mode. 12 | /// 13 | /// Once clients subscribe to a channel, they may only perform pub/sub related 14 | /// commands. The `Client` type is transitioned to a `Subscriber` type in order 15 | /// to prevent non-pub/sub methods from being called. 16 | pub struct Subscriber { 17 | /// The subscribed client. 18 | pub(crate) client: Client, 19 | 20 | /// The set of channels to which the `Subscriber` is currently subscribed. 21 | pub(crate) subscribed_channels: Vec, 22 | } 23 | 24 | /// A message received on a subscribed channel. 25 | #[derive(Debug, Clone)] 26 | pub struct Message { 27 | pub channel: String, 28 | pub content: Bytes, 29 | } 30 | 31 | impl Subscriber { 32 | /// Subscribe to a list of new channels 33 | pub async fn subscribe(&mut self, channels: &[String]) -> Result<(), MiniRedisConnectionError> { 34 | // Issue the subscribe command 35 | self.client.subscribe_cmd(channels).await?; 36 | 37 | // Update the set of subscribed channels. 38 | self.subscribed_channels 39 | .extend(channels.iter().map(Clone::clone)); 40 | 41 | Ok(()) 42 | } 43 | 44 | /// Returns the set of channels currently subscribed to. 45 | pub fn get_subscribed(&self) -> &[String] { 46 | &self.subscribed_channels 47 | } 48 | 49 | /// Receive the next message published on a subscribed channel, waiting if 50 | /// necessary. 51 | /// 52 | /// `None` indicates the subscription has been terminated. 53 | pub async fn next_message(&mut self) -> Result, MiniRedisConnectionError> { 54 | match self.client.connection.read_frame().await? { 55 | Some(frame) => { 56 | debug!("subscribe received next message: {:?}", frame); 57 | 58 | match frame { 59 | Frame::Array(ref frame) => match frame.as_slice() { 60 | [message, channel, content] if *message == "message" => Ok(Some(Message { 61 | channel: channel.to_string(), 62 | content: Bytes::from(content.to_string()), 63 | })), 64 | _ => { 65 | error!("invalid message, frame: {:?}", frame); 66 | Err(MiniRedisConnectionError::InvalidFrameType) 67 | } 68 | }, 69 | frame => Err(MiniRedisConnectionError::CommandExecute(frame.to_string())), 70 | } 71 | } 72 | None => Ok(None), 73 | } 74 | } 75 | 76 | /// Convert the subscriber into a `Stream` yielding new messages published 77 | /// on subscribed channels. 78 | /// 79 | /// `Subscriber` does not implement stream itself as doing so with safe code 80 | /// is non trivial. The usage of async/await would require a manual Stream 81 | /// implementation to use `unsafe` code. Instead, a conversion function is 82 | /// provided and the returned stream is implemented with the help of the 83 | /// `async-stream` crate. 84 | pub fn into_stream(mut self) -> impl Stream> { 85 | // Uses the `try_stream` macro from the `async-stream` crate. Generators 86 | // are not stable in Rust. The crate uses a macro to simulate generators 87 | // on top of async/await. There are limitations, so read the 88 | // documentation there. 89 | try_stream! { 90 | while let Some(message) = self.next_message().await? { 91 | yield message; 92 | } 93 | } 94 | } 95 | 96 | /// Unsubscribe to a list of new channels 97 | pub async fn unsubscribe( 98 | &mut self, 99 | channels: &[String], 100 | ) -> Result<(), MiniRedisConnectionError> { 101 | let frame = Unsubscribe::new(channels).into_frame()?; 102 | 103 | debug!("unsubscribe command: {:?}", frame); 104 | 105 | // Write the frame to the socket 106 | self.client.connection.write_frame(&frame).await?; 107 | 108 | // if the input channel list is empty, server acknowledges as unsubscribing 109 | // from all subscribed channels, so we assert that the unsubscribe list received 110 | // matches the client subscribed one 111 | let num = if channels.is_empty() { 112 | self.subscribed_channels.len() 113 | } else { 114 | channels.len() 115 | }; 116 | 117 | // Read the response 118 | for _ in 0..num { 119 | let response = self.client.read_response().await?; 120 | 121 | match response { 122 | Frame::Array(ref frame) => match frame.as_slice() { 123 | [unsubscribe, channel, ..] if *unsubscribe == "unsubscribe" => { 124 | let len = self.subscribed_channels.len(); 125 | 126 | if len == 0 { 127 | // There must be at least one channel 128 | return Err(MiniRedisConnectionError::InvalidArgument( 129 | response.to_string(), 130 | )); 131 | } 132 | 133 | // unsubscribed channel should exist in the subscribed list at this point 134 | self.subscribed_channels.retain(|c| *channel != &c[..]); 135 | 136 | // Only a single channel should be removed from the 137 | // list of subscribed channels. 138 | if self.subscribed_channels.len() != len - 1 { 139 | return Err(MiniRedisConnectionError::CommandExecute( 140 | response.to_string(), 141 | )); 142 | } 143 | } 144 | _ => { 145 | return Err(MiniRedisConnectionError::InvalidFrameType); 146 | } 147 | }, 148 | frame => return Err(MiniRedisConnectionError::CommandExecute(frame.to_string())), 149 | }; 150 | } 151 | 152 | Ok(()) 153 | } 154 | } 155 | -------------------------------------------------------------------------------- /src/cmd/get.rs: -------------------------------------------------------------------------------- 1 | use bytes::Bytes; 2 | use log::debug; 3 | 4 | use crate::connection::connect::Connection; 5 | use crate::connection::frame::Frame; 6 | use crate::connection::parse::Parse; 7 | use crate::error::{MiniRedisConnectionError, MiniRedisParseError}; 8 | use crate::storage::db::Db; 9 | use crate::storage::traits::KvStore; 10 | 11 | /// Get the value of key. 12 | /// 13 | /// If the key does not exist the special value nil is returned. An error is 14 | /// returned if the value stored at key is not a string, because GET only 15 | /// handles string values. 16 | #[derive(Debug)] 17 | pub struct Get { 18 | /// Name of the key to get 19 | key: String, 20 | } 21 | 22 | impl Get { 23 | /// Create a new `Get` command which fetches `key`. 24 | pub fn new(key: impl ToString) -> Get { 25 | Get { 26 | key: key.to_string(), 27 | } 28 | } 29 | 30 | /// Get the key 31 | pub fn key(&self) -> &str { 32 | &self.key 33 | } 34 | 35 | /// Parse a `Get` instance from a received frame. 36 | /// 37 | /// The `Parse` argument provides a cursor-like API to read fields from the 38 | /// `Frame`. At this point, the entire frame has already been received from 39 | /// the socket. 40 | /// 41 | /// The `GET` string has already been consumed. 42 | /// 43 | /// # Returns 44 | /// 45 | /// Returns the `Get` value on success. If the frame is malformed, `Err` is 46 | /// returned. 47 | /// 48 | /// # Format 49 | /// 50 | /// Expects an array frame containing two entries. 51 | /// 52 | /// ```text 53 | /// GET key 54 | /// ``` 55 | pub(crate) fn parse_frames(parse: &mut Parse) -> Result { 56 | // The `GET` string has already been consumed. The next value is the 57 | // name of the key to get. If the next value is not a string or the 58 | // input is fully consumed, then an error is returned. 59 | let key = parse.next_string()?; 60 | 61 | Ok(Get { key }) 62 | } 63 | 64 | /// Apply the `Get` command to the specified `Db` instance. 65 | /// 66 | /// The response is written to `dst`. This is called by the server in order 67 | /// to execute a received command. 68 | pub(crate) async fn apply( 69 | self, 70 | db: &Db, 71 | dst: &mut Connection, 72 | ) -> Result<(), MiniRedisConnectionError> { 73 | // Get the value from the shared database state 74 | let response = if let Some(value) = db.get(&self.key) { 75 | // If a value is present, it is written to the client in "bulk" format. 76 | Frame::Bulk(value) 77 | } else { 78 | // If there is no value, `Null` is written. 79 | Frame::Null 80 | }; 81 | 82 | debug!("get command applied resp: {:?}", response); 83 | 84 | // Write the response back to the client 85 | dst.write_frame(&response).await?; 86 | 87 | Ok(()) 88 | } 89 | 90 | /// Converts the command into an equivalent `Frame`. 91 | /// 92 | /// This is called by the client when encoding a `Get` command to send to 93 | /// the server. 94 | pub(crate) fn into_frame(self) -> Result { 95 | let mut frame = Frame::array(); 96 | frame.push_bulk(Bytes::from("get".as_bytes()))?; 97 | frame.push_bulk(Bytes::from(self.key.into_bytes()))?; 98 | Ok(frame) 99 | } 100 | } 101 | -------------------------------------------------------------------------------- /src/cmd/mod.rs: -------------------------------------------------------------------------------- 1 | use crate::cmd::get::Get; 2 | use crate::cmd::ping::Ping; 3 | use crate::cmd::publish::Publish; 4 | use crate::cmd::set::Set; 5 | use crate::cmd::subscribe::Subscribe; 6 | use crate::cmd::unknown::Unknown; 7 | use crate::cmd::unsubscribe::Unsubscribe; 8 | use crate::connection::connect::Connection; 9 | use crate::connection::frame::Frame; 10 | use crate::connection::parse::Parse; 11 | use crate::error::{MiniRedisConnectionError, MiniRedisParseError}; 12 | use crate::server::shutdown::Shutdown; 13 | use crate::storage::db::Db; 14 | 15 | pub(crate) mod get; 16 | pub(crate) mod ping; 17 | pub(crate) mod publish; 18 | pub(crate) mod set; 19 | pub(crate) mod subscribe; 20 | pub(crate) mod unknown; 21 | pub(crate) mod unsubscribe; 22 | 23 | /// Enumeration of supported Redis commands. 24 | /// 25 | /// Methods called on `Command` are delegated to the command implementation. 26 | #[derive(Debug)] 27 | pub enum Command { 28 | Get(Get), 29 | Set(Set), 30 | Publish(Publish), 31 | Subscribe(Subscribe), 32 | Unsubscribe(Unsubscribe), 33 | Ping(Ping), 34 | Unknown(Unknown), 35 | } 36 | 37 | impl Command { 38 | /// Parse a command from a received frame. 39 | /// 40 | /// The `Frame` must represent a Redis command supported by `mini-redis` and 41 | /// be the array variant. 42 | /// 43 | /// # Returns 44 | /// 45 | /// On success, the command value is returned, otherwise, `Err` is returned. 46 | pub fn from_frame(frame: Frame) -> Result { 47 | // The frame value is decorated with `Parse`. `Parse` provides a 48 | // "cursor" like API which makes parsing the command easier. 49 | // 50 | // The frame value must be an array variant. Any other frame variants 51 | // result in an error being returned. 52 | let mut parse = Parse::new(frame)?; 53 | 54 | // All redis commands begin with the command name as a string. The name 55 | // is read and converted to lower cases in order to do case sensitive 56 | // matching. 57 | let command_name = parse.next_string()?.to_lowercase(); 58 | 59 | // Match the command name, delegating the rest of the parsing to the 60 | // specific command. 61 | let command = match &command_name[..] { 62 | "get" => Command::Get(Get::parse_frames(&mut parse)?), 63 | "set" => Command::Set(Set::parse_frames(&mut parse)?), 64 | "publish" => Command::Publish(Publish::parse_frames(&mut parse)?), 65 | "subscribe" => Command::Subscribe(Subscribe::parse_frames(&mut parse)?), 66 | "unsubscribe" => Command::Unsubscribe(Unsubscribe::parse_frames(&mut parse)?), 67 | "ping" => Command::Ping(Ping::parse_frames(&mut parse)?), 68 | _ => { 69 | // The command is not recognized and an Unknown command is 70 | // returned. 71 | // 72 | // `return` is called here to skip the `finish()` call below. As 73 | // the command is not recognized, there is most likely 74 | // unconsumed fields remaining in the `Parse` instance. 75 | return Ok(Command::Unknown(Unknown::new(command_name))); 76 | } 77 | }; 78 | 79 | // Check if there is any remaining unconsumed fields in the `Parse` 80 | // value. If fields remain, this indicates an unexpected frame format 81 | // and an error is returned. 82 | parse.finish()?; 83 | 84 | // The command has been successfully parsed 85 | Ok(command) 86 | } 87 | 88 | /// Apply the command to the specified `Db` instance. 89 | /// 90 | /// The response is written to `dst`. This is called by the server in order 91 | /// to execute a received command. 92 | /// Apply the command to the specified `Db` instance. 93 | /// 94 | /// The response is written to `dst`. This is called by the server in order 95 | /// to execute a received command. 96 | pub(crate) async fn apply( 97 | self, 98 | db: &Db, 99 | dst: &mut Connection, 100 | shutdown: &mut Shutdown, 101 | ) -> Result<(), MiniRedisConnectionError> { 102 | use Command::*; 103 | 104 | match self { 105 | Ping(cmd) => cmd.apply(dst).await, 106 | Get(cmd) => cmd.apply(db, dst).await, 107 | Set(cmd) => cmd.apply(db, dst).await, 108 | Publish(cmd) => cmd.apply(db, dst).await, 109 | Subscribe(cmd) => cmd.apply(db, dst, shutdown).await, 110 | // `Unsubscribe` cannot be applied. It may only be received from the 111 | // context of a `Subscribe` command. 112 | Unsubscribe(_) => Err(MiniRedisConnectionError::CommandExecute( 113 | "`Unsubscribe` is unsupported in this context".into(), 114 | )), 115 | Unknown(cmd) => cmd.apply(dst).await, 116 | } 117 | } 118 | 119 | /// Returns the command name 120 | pub(crate) fn get_name(&self) -> &str { 121 | match self { 122 | Command::Get(_) => "get", 123 | Command::Set(_) => "set", 124 | Command::Publish(_) => "pub", 125 | Command::Subscribe(_) => "subscribe", 126 | Command::Unsubscribe(_) => "unsubscribe", 127 | Command::Ping(_) => "ping", 128 | Command::Unknown(cmd) => cmd.get_name(), 129 | } 130 | } 131 | } 132 | -------------------------------------------------------------------------------- /src/cmd/ping.rs: -------------------------------------------------------------------------------- 1 | use crate::connection::connect::Connection; 2 | use crate::connection::frame::Frame; 3 | use crate::connection::parse::Parse; 4 | use crate::error::{MiniRedisConnectionError, MiniRedisParseError}; 5 | use bytes::Bytes; 6 | 7 | /// Returns PONG if no argument is provided, otherwise 8 | /// return a copy of the argument as a bulk. 9 | /// 10 | /// This command is often used to test if a connection 11 | /// is still alive, or to measure latency. 12 | #[derive(Debug, Default)] 13 | pub struct Ping { 14 | /// optional message to be returned 15 | msg: Option, 16 | } 17 | 18 | impl Ping { 19 | /// Create a new `Ping` command with optional `msg`. 20 | pub fn new(msg: Option) -> Ping { 21 | Ping { msg } 22 | } 23 | 24 | /// Parse a `Ping` instance from a received frame. 25 | /// 26 | /// The `Parse` argument provides a cursor-like API to read fields from the 27 | /// `Frame`. At this point, the entire frame has already been received from 28 | /// the socket. 29 | /// 30 | /// The `PING` string has already been consumed. 31 | /// 32 | /// # Returns 33 | /// 34 | /// Returns the `Ping` value on success. If the frame is malformed, `Err` is 35 | /// returned. 36 | /// 37 | /// # Format 38 | /// 39 | /// Expects an array frame containing `PING` and an optional message. 40 | /// 41 | /// ```text 42 | /// PING [message] 43 | /// ``` 44 | pub(crate) fn parse_frames(parse: &mut Parse) -> Result { 45 | match parse.next_string() { 46 | Ok(msg) => Ok(Ping::new(Some(msg))), 47 | Err(MiniRedisParseError::EndOfStream) => Ok(Ping::default()), 48 | Err(e) => Err(e), 49 | } 50 | } 51 | 52 | /// Apply the `Ping` command and return the message. 53 | /// 54 | /// The response is written to `dst`. This is called by the server in order 55 | /// to execute a received command. 56 | pub(crate) async fn apply(self, dst: &mut Connection) -> Result<(), MiniRedisConnectionError> { 57 | let response = match self.msg { 58 | None => Frame::Simple("PONG".to_string()), 59 | Some(msg) => Frame::Bulk(Bytes::from(msg)), 60 | }; 61 | 62 | // Write the response back to the client 63 | dst.write_frame(&response).await?; 64 | 65 | Ok(()) 66 | } 67 | 68 | /// Converts the command into an equivalent `Frame`. 69 | /// 70 | /// This is called by the client when encoding a `Ping` command to send 71 | /// to the server. 72 | pub(crate) fn into_frame(self) -> Result { 73 | let mut frame = Frame::array(); 74 | frame.push_bulk(Bytes::from("ping".as_bytes()))?; 75 | if let Some(msg) = self.msg { 76 | frame.push_bulk(Bytes::from(msg))?; 77 | } 78 | Ok(frame) 79 | } 80 | } 81 | -------------------------------------------------------------------------------- /src/cmd/publish.rs: -------------------------------------------------------------------------------- 1 | use crate::connection::connect::Connection; 2 | use crate::connection::frame::Frame; 3 | use crate::connection::parse::Parse; 4 | use crate::error::{MiniRedisConnectionError, MiniRedisParseError}; 5 | use crate::storage::db::Db; 6 | use crate::storage::traits::KvStore; 7 | use bytes::Bytes; 8 | use log::debug; 9 | 10 | /// Posts a message to the given channel. 11 | /// 12 | /// Send a message into a channel without any knowledge of individual consumers. 13 | /// Consumers may subscribe to channels in order to receive the messages. 14 | /// 15 | /// Channel names have no relation to the key-value namespace. Publishing on a 16 | /// channel named "foo" has no relation to setting the "foo" key. 17 | #[derive(Debug)] 18 | pub struct Publish { 19 | /// Name of the channel on which the message should be published. 20 | channel: String, 21 | 22 | /// The message to publish. 23 | message: Bytes, 24 | } 25 | 26 | impl Publish { 27 | /// Create a new `Publish` command which sends `message` on `channel`. 28 | pub(crate) fn new(channel: impl ToString, message: Bytes) -> Self { 29 | Publish { 30 | channel: channel.to_string(), 31 | message, 32 | } 33 | } 34 | 35 | /// Parse a `Publish` instance from a received frame. 36 | /// 37 | /// The `Parse` argument provides a cursor-like API to read fields from the 38 | /// `Frame`. At this point, the entire frame has already been received from 39 | /// the socket. 40 | /// 41 | /// The `PUBLISH` string has already been consumed. 42 | /// 43 | /// # Returns 44 | /// 45 | /// On success, the `Publish` value is returned. If the frame is malformed, 46 | /// `Err` is returned. 47 | /// 48 | /// # Format 49 | /// 50 | /// Expects an array frame containing three entries. 51 | /// 52 | /// ```text 53 | /// PUBLISH channel message 54 | /// ``` 55 | pub(crate) fn parse_frames(parse: &mut Parse) -> Result { 56 | // The `PUBLISH` string has already been consumed. Extract the `channel` 57 | // and `message` values from the frame. 58 | // 59 | // The `channel` must be a valid string. 60 | let channel = parse.next_string()?; 61 | 62 | // The `message` is arbitrary bytes. 63 | let message = parse.next_bytes()?; 64 | 65 | Ok(Publish { channel, message }) 66 | } 67 | 68 | /// Apply the `Publish` command to the specified `Db` instance. 69 | /// 70 | /// The response is written to `dst`. This is called by the server in order 71 | /// to execute a received command. 72 | pub(crate) async fn apply( 73 | self, 74 | db: &Db, 75 | dst: &mut Connection, 76 | ) -> Result<(), MiniRedisConnectionError> { 77 | // The shared state contains the `tokio::sync::broadcast::Sender` for 78 | // all active channels. Calling `db.publish` dispatches the message into 79 | // the appropriate channel. 80 | // 81 | // The number of subscribers currently listening on the channel is 82 | // returned. This does not mean that `num_subscriber` channels will 83 | // receive the message. Subscribers may drop before receiving the 84 | // message. Given this, `num_subscribers` should only be used as a 85 | // "hint". 86 | let num_subscribers = db.publish(&self.channel, self.message); 87 | 88 | // The number of subscribers is returned as the response to the publish 89 | // request. 90 | let response = Frame::Integer(num_subscribers as u64); 91 | debug!("apply command applied response: {}", response); 92 | 93 | // Write the frame to the client. 94 | dst.write_frame(&response).await?; 95 | 96 | Ok(()) 97 | } 98 | 99 | /// Converts the command into an equivalent `Frame`. 100 | /// 101 | /// This is called by the client when encoding a `Publish` command to send 102 | /// to the server. 103 | pub(crate) fn into_frame(self) -> Result { 104 | let mut frame = Frame::array(); 105 | frame.push_bulk(Bytes::from("publish".as_bytes()))?; 106 | frame.push_bulk(Bytes::from(self.channel.into_bytes()))?; 107 | frame.push_bulk(self.message)?; 108 | 109 | Ok(frame) 110 | } 111 | } 112 | -------------------------------------------------------------------------------- /src/cmd/set.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use bytes::Bytes; 4 | use log::{debug, warn}; 5 | 6 | use crate::connection::connect::Connection; 7 | use crate::connection::frame::Frame; 8 | use crate::connection::parse::Parse; 9 | use crate::error::{MiniRedisConnectionError, MiniRedisParseError}; 10 | use crate::storage::db::Db; 11 | use crate::storage::traits::KvStore; 12 | 13 | /// Set `key` to hold the string `value`. 14 | /// 15 | /// If `key` already holds a value, it is overwritten, regardless of its type. 16 | /// Any previous time to live associated with the key is discarded on successful 17 | /// SET operation. 18 | /// 19 | /// # Options 20 | /// 21 | /// Currently, the following options are supported: 22 | /// 23 | /// * EX `seconds` -- Set the specified expire time, in seconds. 24 | /// * PX `milliseconds` -- Set the specified expire time, in milliseconds. 25 | #[derive(Debug)] 26 | pub struct Set { 27 | /// the lookup key 28 | key: String, 29 | 30 | /// the value to be stored 31 | value: Bytes, 32 | 33 | /// When to expire the key 34 | expire: Option, 35 | } 36 | 37 | impl Set { 38 | /// Create a new `Set` command which sets `key` to `value`. 39 | /// 40 | /// If `expire` is `Some`, the value should expire after the specified 41 | /// duration. 42 | pub fn new(key: impl ToString, value: Bytes, expire: Option) -> Set { 43 | Set { 44 | key: key.to_string(), 45 | value, 46 | expire, 47 | } 48 | } 49 | 50 | /// Parse a `Set` instance from a received frame. 51 | /// 52 | /// The `Parse` argument provides a cursor-like API to read fields from the 53 | /// `Frame`. At this point, the entire frame has already been received from 54 | /// the socket. 55 | /// 56 | /// The `SET` string has already been consumed. 57 | /// 58 | /// # Returns 59 | /// 60 | /// Returns the `Set` value on success. If the frame is malformed, `Err` is 61 | /// returned. 62 | /// 63 | /// # Format 64 | /// 65 | /// Expects an array frame containing at least 3 entries. 66 | /// 67 | /// ```text 68 | /// SET key value [EX seconds|PX milliseconds] 69 | /// ``` 70 | pub(crate) fn parse_frames(parse: &mut Parse) -> Result { 71 | // Read the key to set. This is a required field 72 | let key = parse.next_string()?; 73 | 74 | // Read the value to set. This is a required field. 75 | let value = parse.next_bytes()?; 76 | 77 | // The expiration is optional. If nothing else follows, then it is `None`. 78 | let mut expire = None; 79 | 80 | // Attempt to parse another string. 81 | match parse.next_string() { 82 | Ok(s) if s.to_uppercase() == "EX" => { 83 | // An expiration is specified in seconds. The next value is an 84 | // integer. 85 | let secs = parse.next_int()?; 86 | expire = Some(Duration::from_secs(secs)); 87 | } 88 | Ok(s) if s.to_uppercase() == "PX" => { 89 | // An expiration is specified in milliseconds. The next value is 90 | // an integer. 91 | let ms = parse.next_int()?; 92 | expire = Some(Duration::from_millis(ms)); 93 | } 94 | // Currently, mini-redis does not support any of the other SET 95 | // options. An error here results in the connection being 96 | // terminated. Other connections will continue to operate normally. 97 | Ok(s) => { 98 | warn!("unsupported SET option: {}", s); 99 | return Err(MiniRedisParseError::Parse( 100 | "currently `SET` only supports the expiration option".into(), 101 | )); 102 | } 103 | // The `EndOfStream` error indicates there is no further data to 104 | // parse. In this case, it is a normal run time situation and 105 | // indicates there are no specified `SET` options. 106 | Err(MiniRedisParseError::EndOfStream) => { 107 | debug!("no extra SET option"); 108 | } 109 | // All other errors are bubbled up, resulting in the connection 110 | // being terminated. 111 | Err(err) => return Err(err), 112 | } 113 | 114 | Ok(Set { key, value, expire }) 115 | } 116 | 117 | /// Apply the `Set` command to the specified `Db` instance. 118 | /// 119 | /// The response is written to `dst`. This is called by the server in order 120 | /// to execute a received command. 121 | pub(crate) async fn apply( 122 | self, 123 | db: &Db, 124 | dst: &mut Connection, 125 | ) -> Result<(), MiniRedisConnectionError> { 126 | // Set the value in the shared database state. 127 | db.set(self.key, self.value, self.expire); 128 | 129 | // Create a success response and write it to `dst`. 130 | let response = Frame::Simple("OK".to_string()); 131 | debug!("applied set command response: {:?}", response); 132 | 133 | dst.write_frame(&response).await?; 134 | 135 | Ok(()) 136 | } 137 | 138 | /// Converts the command into an equivalent `Frame`. 139 | /// 140 | /// This is called by the client when encoding a `Set` command to send to 141 | /// the server. 142 | pub(crate) fn into_frame(self) -> Result { 143 | let mut frame = Frame::array(); 144 | frame.push_bulk(Bytes::from("set".as_bytes()))?; 145 | frame.push_bulk(Bytes::from(self.key.into_bytes()))?; 146 | frame.push_bulk(self.value)?; 147 | if let Some(ms) = self.expire { 148 | // Expirations in Redis procotol can be specified in two ways 149 | // 1. SET key value EX seconds 150 | // 2. SET key value PX milliseconds 151 | // We implement the second option because it allows greater precision and 152 | // src/bin/cli.rs parses the expiration argument as milliseconds 153 | // in duration_from_ms_str() 154 | frame.push_bulk(Bytes::from("px".as_bytes()))?; 155 | frame.push_int(ms.as_millis() as u64)?; 156 | } 157 | Ok(frame) 158 | } 159 | 160 | /// Get the key 161 | pub fn key(&self) -> &str { 162 | &self.key 163 | } 164 | 165 | /// Get the value 166 | pub fn value(&self) -> &Bytes { 167 | &self.value 168 | } 169 | 170 | /// Get the expire 171 | pub fn expire(&self) -> Option { 172 | self.expire 173 | } 174 | } 175 | -------------------------------------------------------------------------------- /src/cmd/subscribe.rs: -------------------------------------------------------------------------------- 1 | use std::pin::Pin; 2 | 3 | use bytes::Bytes; 4 | use log::{debug, warn}; 5 | use tokio::select; 6 | use tokio_stream::{Stream, StreamExt, StreamMap}; 7 | 8 | use crate::cmd::unknown::Unknown; 9 | use crate::cmd::unsubscribe::make_unsubscribe_frame; 10 | use crate::cmd::Command; 11 | use crate::connection::connect::Connection; 12 | use crate::connection::frame::Frame; 13 | use crate::connection::parse::Parse; 14 | use crate::error::{MiniRedisConnectionError, MiniRedisParseError}; 15 | use crate::server::shutdown::Shutdown; 16 | use crate::storage::db::Db; 17 | use crate::storage::traits::KvStore; 18 | 19 | /// Subscribes the client to one or more channels. 20 | /// 21 | /// Once the client enters the subscribed state, it is not supposed to issue any 22 | /// other commands, except for additional SUBSCRIBE, PSUBSCRIBE, UNSUBSCRIBE, 23 | /// PUNSUBSCRIBE, PING and QUIT commands. 24 | #[derive(Debug)] 25 | pub struct Subscribe { 26 | channels: Vec, 27 | } 28 | 29 | /// Stream of messages. The stream receives messages from the 30 | /// `broadcast::Receiver`. We use `stream!` to create a `Stream` that consumes 31 | /// messages. Because `stream!` values cannot be named, we box the stream using 32 | /// a trait object. 33 | type Messages = Pin + Send>>; 34 | 35 | impl Subscribe { 36 | /// Creates a new `Subscribe` command to listen on the specified channels. 37 | pub(crate) fn new(channels: &[String]) -> Subscribe { 38 | Subscribe { 39 | channels: channels.to_vec(), 40 | } 41 | } 42 | 43 | /// Parse a `Subscribe` instance from a received frame. 44 | /// 45 | /// The `Parse` argument provides a cursor-like API to read fields from the 46 | /// `Frame`. At this point, the entire frame has already been received from 47 | /// the socket. 48 | /// 49 | /// The `SUBSCRIBE` string has already been consumed. 50 | /// 51 | /// # Returns 52 | /// 53 | /// On success, the `Subscribe` value is returned. If the frame is 54 | /// malformed, `Err` is returned. 55 | /// 56 | /// # Format 57 | /// 58 | /// Expects an array frame containing two or more entries. 59 | /// 60 | /// ```text 61 | /// SUBSCRIBE channel [channel ...] 62 | /// ``` 63 | pub(crate) fn parse_frames(parse: &mut Parse) -> Result { 64 | // The `SUBSCRIBE` string has already been consumed. At this point, 65 | // there is one or more strings remaining in `parse`. These represent 66 | // the channels to subscribe to. 67 | // 68 | // Extract the first string. If there is none, the the frame is 69 | // malformed and the error is bubbled up. 70 | let mut channels = vec![parse.next_string()?]; 71 | 72 | // Now, the remainder of the frame is consumed. Each value must be a 73 | // string or the frame is malformed. Once all values in the frame have 74 | // been consumed, the command is fully parsed. 75 | loop { 76 | match parse.next_string() { 77 | // A string has been consumed from the `parse`, push it into the 78 | // list of channels to subscribe to. 79 | Ok(s) => channels.push(s), 80 | // The `EndOfStream` error indicates there is no further data to 81 | // parse. 82 | Err(MiniRedisParseError::EndOfStream) => break, 83 | // All other errors are bubbled up, resulting in the connection 84 | // being terminated. 85 | Err(err) => return Err(err), 86 | } 87 | } 88 | 89 | Ok(Subscribe { channels }) 90 | } 91 | 92 | /// Apply the `Subscribe` command to the specified `Db` instance. 93 | /// 94 | /// This function is the entry point and includes the initial list of 95 | /// channels to subscribe to. Additional `subscribe` and `unsubscribe` 96 | /// commands may be received from the client and the list of subscriptions 97 | /// are updated accordingly. 98 | /// 99 | /// [here]: https://redis.io/topics/pubsub 100 | pub(crate) async fn apply( 101 | mut self, 102 | db: &Db, 103 | dst: &mut Connection, 104 | shutdown: &mut Shutdown, 105 | ) -> Result<(), MiniRedisConnectionError> { 106 | // Each individual channel subscription is handled using a 107 | // `sync::broadcast` channel. Messages are then fanned out to all 108 | // clients currently subscribed to the channels. 109 | // 110 | // An individual client may subscribe to multiple channels and may 111 | // dynamically add and remove channels from its subscription set. To 112 | // handle this, a `StreamMap` is used to track active subscriptions. The 113 | // `StreamMap` merges messages from individual broadcast channels as 114 | // they are received. 115 | let mut subscriptions = StreamMap::new(); 116 | 117 | loop { 118 | // `self.channels` is used to track additional channels to subscribe 119 | // to. When new `SUBSCRIBE` commands are received during the 120 | // execution of `apply`, the new channels are pushed onto this vec. 121 | for channel_name in self.channels.drain(..) { 122 | Self::subscribe_to_channel(channel_name, &mut subscriptions, db, dst).await?; 123 | } 124 | 125 | // Wait for one of the following to happen: 126 | // 127 | // - Receive a message from one of the subscribed channels. 128 | // - Receive a subscribe or unsubscribe command from the client. 129 | // - A server shutdown signal. 130 | select! { 131 | // Receive messages from subscribed channels 132 | Some((channel_name, msg)) = subscriptions.next() => { 133 | dst.write_frame(&make_message_frame(channel_name, msg)?).await?; 134 | } 135 | res = dst.read_frame() => { 136 | let frame = match res? { 137 | Some(frame) => frame, 138 | // This happens if the remote client has disconnected. 139 | None => { 140 | warn!("remote subscribe client disconnected"); 141 | return Ok(()) 142 | } 143 | }; 144 | 145 | handle_command( 146 | frame, 147 | &mut self.channels, 148 | &mut subscriptions, 149 | dst, 150 | ).await?; 151 | } 152 | _ = shutdown.recv() => { 153 | warn!("server shutdown, stop subscribe"); 154 | return Ok(()); 155 | } 156 | } 157 | } 158 | } 159 | 160 | /// Converts the command into an equivalent `Frame`. 161 | /// 162 | /// This is called by the client when encoding a `Subscribe` command to send 163 | /// to the server. 164 | pub(crate) fn into_frame(self) -> Result { 165 | let mut frame = Frame::array(); 166 | frame.push_bulk(Bytes::from("subscribe".as_bytes()))?; 167 | for channel in self.channels { 168 | frame.push_bulk(Bytes::from(channel.into_bytes()))?; 169 | } 170 | Ok(frame) 171 | } 172 | 173 | async fn subscribe_to_channel( 174 | channel_name: String, 175 | subscriptions: &mut StreamMap, 176 | db: &Db, 177 | dst: &mut Connection, 178 | ) -> Result<(), MiniRedisConnectionError> { 179 | let mut rx = db.subscribe(channel_name.clone()); 180 | 181 | // Subscribe to the channel. 182 | let rx = Box::pin(async_stream::stream! { 183 | loop { 184 | match rx.recv().await { 185 | Ok(msg) => yield msg, 186 | // If we lagged in consuming messages, just resume. 187 | Err(tokio::sync::broadcast::error::RecvError::Lagged(e)) => { 188 | warn!("subscribe received lagged: {}", e); 189 | } 190 | Err(e) => { 191 | warn!("subscribe received error: {}", e); 192 | break 193 | }, 194 | } 195 | } 196 | }); 197 | 198 | // Track subscription in this client's subscription set. 199 | subscriptions.insert(channel_name.clone(), rx); 200 | 201 | debug!("subscribed to channel success: {}", channel_name); 202 | 203 | // Respond with the successful subscription 204 | let response = make_subscribe_frame(channel_name, subscriptions.len())?; 205 | dst.write_frame(&response).await?; 206 | 207 | Ok(()) 208 | } 209 | } 210 | 211 | /// Creates the response to a subscribe request. 212 | /// 213 | /// All of these functions take the `channel_name` as a `String` instead of 214 | /// a `&str` since `Bytes::from` can reuse the allocation in the `String`, and 215 | /// taking a `&str` would require copying the data. This allows the caller to 216 | /// decide whether to clone the channel name or not. 217 | fn make_subscribe_frame( 218 | channel_name: String, 219 | num_subs: usize, 220 | ) -> Result { 221 | let mut response = Frame::array(); 222 | response.push_bulk(Bytes::from_static(b"subscribe"))?; 223 | response.push_bulk(Bytes::from(channel_name))?; 224 | response.push_int(num_subs as u64)?; 225 | Ok(response) 226 | } 227 | 228 | /// Creates a message informing the client about a new message on a channel that 229 | /// the client subscribes to. 230 | fn make_message_frame(channel_name: String, msg: Bytes) -> Result { 231 | let mut response = Frame::array(); 232 | response.push_bulk(Bytes::from_static(b"message"))?; 233 | response.push_bulk(Bytes::from(channel_name))?; 234 | response.push_bulk(msg)?; 235 | Ok(response) 236 | } 237 | 238 | /// Handle a command received while inside `Subscribe::apply`. Only subscribe 239 | /// and unsubscribe commands are permitted in this context. 240 | /// 241 | /// Any new subscriptions are appended to `subscribe_to` instead of modifying 242 | /// `subscriptions`. 243 | async fn handle_command( 244 | frame: Frame, 245 | subscribe_to: &mut Vec, 246 | subscriptions: &mut StreamMap, 247 | dst: &mut Connection, 248 | ) -> Result<(), MiniRedisConnectionError> { 249 | // A command has been received from the client. 250 | // 251 | // Only `SUBSCRIBE` and `UNSUBSCRIBE` commands are permitted 252 | // in this context. 253 | match Command::from_frame(frame)? { 254 | Command::Subscribe(subscribe) => { 255 | // The `apply` method will subscribe to the channels we add to this 256 | // vector. 257 | subscribe_to.extend(subscribe.channels.into_iter()); 258 | } 259 | Command::Unsubscribe(mut unsubscribe) => { 260 | // If no channels are specified, this requests unsubscribing from 261 | // **all** channels. To implement this, the `unsubscribe.channels` 262 | // vec is populated with the list of channels currently subscribed 263 | // to. 264 | if unsubscribe.channels.is_empty() { 265 | unsubscribe.channels = subscriptions 266 | .keys() 267 | .map(|channel_name| channel_name.to_string()) 268 | .collect(); 269 | } 270 | 271 | for channel_name in unsubscribe.channels { 272 | debug!("begin unsubscribed: {}", channel_name); 273 | subscriptions.remove(&channel_name); 274 | 275 | let response = make_unsubscribe_frame(channel_name, subscriptions.len())?; 276 | dst.write_frame(&response).await?; 277 | debug!("unsubscribed success: {}", response); 278 | } 279 | } 280 | command => { 281 | let cmd = Unknown::new(command.get_name()); 282 | cmd.apply(dst).await?; 283 | } 284 | } 285 | Ok(()) 286 | } 287 | -------------------------------------------------------------------------------- /src/cmd/unknown.rs: -------------------------------------------------------------------------------- 1 | use log::debug; 2 | 3 | use crate::connection::connect::Connection; 4 | use crate::connection::frame::Frame; 5 | use crate::error::MiniRedisConnectionError; 6 | 7 | /// Represents an "unknown" command. This is not a real `Redis` command. 8 | #[derive(Debug)] 9 | pub struct Unknown { 10 | command_name: String, 11 | } 12 | 13 | impl Unknown { 14 | /// Create a new `Unknown` command which responds to unknown commands 15 | /// issued by clients 16 | pub(crate) fn new(key: impl ToString) -> Unknown { 17 | Unknown { 18 | command_name: key.to_string(), 19 | } 20 | } 21 | 22 | /// Returns the command name 23 | pub(crate) fn get_name(&self) -> &str { 24 | &self.command_name 25 | } 26 | 27 | /// Responds to the client, indicating the command is not recognized. 28 | /// 29 | /// This usually means the command is not yet implemented by `mini-redis`. 30 | pub(crate) async fn apply(self, dst: &mut Connection) -> Result<(), MiniRedisConnectionError> { 31 | let response = Frame::Error(format!("err unknown command '{}'", self.command_name)); 32 | 33 | debug!("apply unknown command resp: {:?}", response); 34 | 35 | dst.write_frame(&response).await?; 36 | Ok(()) 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /src/cmd/unsubscribe.rs: -------------------------------------------------------------------------------- 1 | use bytes::Bytes; 2 | 3 | use crate::connection::frame::Frame; 4 | use crate::connection::parse::Parse; 5 | use crate::error::MiniRedisParseError; 6 | 7 | /// Unsubscribes the client from one or more channels. 8 | /// 9 | /// When no channels are specified, the client is unsubscribed from all the 10 | /// previously subscribed channels. 11 | #[derive(Clone, Debug)] 12 | pub struct Unsubscribe { 13 | pub(crate) channels: Vec, 14 | } 15 | 16 | impl Unsubscribe { 17 | /// Create a new `Unsubscribe` command with the given `channels`. 18 | pub(crate) fn new(channels: &[String]) -> Unsubscribe { 19 | Unsubscribe { 20 | channels: channels.to_vec(), 21 | } 22 | } 23 | 24 | /// Parse a `Unsubscribe` instance from a received frame. 25 | /// 26 | /// The `Parse` argument provides a cursor-like API to read fields from the 27 | /// `Frame`. At this point, the entire frame has already been received from 28 | /// the socket. 29 | /// 30 | /// The `UNSUBSCRIBE` string has already been consumed. 31 | /// 32 | /// # Returns 33 | /// 34 | /// On success, the `Unsubscribe` value is returned. If the frame is 35 | /// malformed, `Err` is returned. 36 | /// 37 | /// # Format 38 | /// 39 | /// Expects an array frame containing at least one entry. 40 | /// 41 | /// ```text 42 | /// UNSUBSCRIBE [channel [channel ...]] 43 | /// ``` 44 | pub(crate) fn parse_frames(parse: &mut Parse) -> Result { 45 | // There may be no channels listed, so start with an empty vec. 46 | let mut channels = vec![]; 47 | 48 | // Each entry in the frame must be a string or the frame is malformed. 49 | // Once all values in the frame have been consumed, the command is fully 50 | // parsed. 51 | loop { 52 | match parse.next_string() { 53 | // A string has been consumed from the `parse`, push it into the 54 | // list of channels to unsubscribe from. 55 | Ok(s) => channels.push(s), 56 | // The `EndOfStream` error indicates there is no further data to 57 | // parse. 58 | Err(MiniRedisParseError::EndOfStream) => break, 59 | // All other errors are bubbled up, resulting in the connection 60 | // being terminated. 61 | Err(err) => return Err(err), 62 | } 63 | } 64 | 65 | Ok(Unsubscribe { channels }) 66 | } 67 | 68 | /// Converts the command into an equivalent `Frame`. 69 | /// 70 | /// This is called by the client when encoding an `Unsubscribe` command to 71 | /// send to the server. 72 | pub(crate) fn into_frame(self) -> Result { 73 | let mut frame = Frame::array(); 74 | frame.push_bulk(Bytes::from("unsubscribe".as_bytes()))?; 75 | 76 | for channel in self.channels { 77 | frame.push_bulk(Bytes::from(channel.into_bytes()))?; 78 | } 79 | 80 | Ok(frame) 81 | } 82 | } 83 | 84 | /// Creates the response to an unsubscribe request. 85 | pub(crate) fn make_unsubscribe_frame( 86 | channel_name: String, 87 | num_subs: usize, 88 | ) -> Result { 89 | let mut response = Frame::array(); 90 | response.push_bulk(Bytes::from_static(b"unsubscribe"))?; 91 | response.push_bulk(Bytes::from(channel_name))?; 92 | response.push_int(num_subs as u64)?; 93 | Ok(response) 94 | } 95 | -------------------------------------------------------------------------------- /src/config.rs: -------------------------------------------------------------------------------- 1 | /// Logger level 2 | pub static LOG_LEVEL: &str = "LOG_LEVEL"; 3 | -------------------------------------------------------------------------------- /src/connection/connect.rs: -------------------------------------------------------------------------------- 1 | use std::io::Cursor; 2 | 3 | use bytes::{Buf, BytesMut}; 4 | use log::warn; 5 | use tokio::io::{AsyncReadExt, AsyncWriteExt, BufWriter}; 6 | use tokio::net::TcpStream; 7 | 8 | use crate::connection::frame::Frame; 9 | use crate::error::{MiniRedisConnectionError, MiniRedisParseError}; 10 | 11 | /// Send and receive `Frame` values from a remote peer. 12 | /// 13 | /// When implementing networking protocols, a message on that protocol is 14 | /// often composed of several smaller messages known as frames. The purpose of 15 | /// `Connection` is to read and write frames on the underlying `TcpStream`. 16 | /// 17 | /// To read frames, the `Connection` uses an internal buffer, which is filled 18 | /// up until there are enough bytes to create a full frame. Once this happens, 19 | /// the `Connection` creates the frame and returns it to the caller. 20 | /// 21 | /// When sending frames, the frame is first encoded into the write buffer. 22 | /// The contents of the write buffer are then written to the socket. 23 | #[derive(Debug)] 24 | pub struct Connection { 25 | /// The `TcpStream`. It is decorated with a `BufWriter`, which provides write 26 | /// level buffering. The `BufWriter` implementation provided by Tokio is 27 | /// sufficient for our needs. 28 | stream: BufWriter, 29 | 30 | // The buffer for reading frames. 31 | buffer: BytesMut, 32 | } 33 | 34 | impl Connection { 35 | /// Create a new `Connection`, backed by `socket`. Read and write buffers 36 | /// are initialized. 37 | pub fn new(socket: TcpStream) -> Connection { 38 | Connection { 39 | stream: BufWriter::new(socket), 40 | // Default to a 4KB read buffer. For the use case of mini redis, 41 | // this is fine. However, real applications will want to tune this 42 | // value to their specific use case. There is a high likelihood that 43 | // a larger read buffer will work better. 44 | buffer: BytesMut::with_capacity(4 * 1024), 45 | } 46 | } 47 | 48 | /// Read a single `Frame` value from the underlying stream. 49 | /// 50 | /// The function waits until it has retrieved enough data to parse a frame. 51 | /// Any data remaining in the read buffer after the frame has been parsed is 52 | /// kept there for the next call to `read_frame`. 53 | /// 54 | /// # Returns 55 | /// 56 | /// On success, the received frame is returned. If the `TcpStream` 57 | /// is closed in a way that doesn't break a frame in half, it returns 58 | /// `None`. Otherwise, an error is returned. 59 | pub async fn read_frame(&mut self) -> Result, MiniRedisConnectionError> { 60 | loop { 61 | // Attempt to parse a frame from the buffered data. If enough data 62 | // has been buffered, the frame is returned. 63 | if let Some(frame) = self.parse_frame()? { 64 | return Ok(Some(frame)); 65 | } 66 | 67 | // There is not enough buffered data to read a frame. Attempt to 68 | // read more data from the socket. 69 | // 70 | // On success, the number of bytes is returned. `0` indicates "end 71 | // of stream". 72 | if 0 == self.stream.read_buf(&mut self.buffer).await? { 73 | // The remote closed the connection. For this to be a clean 74 | // shutdown, there should be no data in the read buffer. If 75 | // there is, this means that the peer closed the socket while 76 | // sending a frame. 77 | return if self.buffer.is_empty() { 78 | Ok(None) 79 | } else { 80 | Err(MiniRedisConnectionError::Disconnect) 81 | }; 82 | } 83 | } 84 | } 85 | 86 | /// Tries to parse a frame from the buffer. If the buffer contains enough 87 | /// data, the frame is returned and the data removed from the buffer. If not 88 | /// enough data has been buffered yet, `Ok(None)` is returned. If the 89 | /// buffered data does not represent a valid frame, `Err` is returned. 90 | fn parse_frame(&mut self) -> Result, MiniRedisConnectionError> { 91 | // Cursor is used to track the "current" location in the 92 | // buffer. Cursor also implements `Buf` from the `bytes` crate 93 | // which provides a number of helpful utilities for working 94 | // with bytes. 95 | let mut buf = Cursor::new(&self.buffer[..]); 96 | 97 | // The first step is to check if enough data has been buffered to parse a single frame. 98 | // This step is usually much faster than doing a full 99 | // parse of the frame, and allows us to skip allocating data structures 100 | // to hold the frame data unless we know the full frame has been received. 101 | match Frame::check(&mut buf) { 102 | Ok(_) => { 103 | // The `check` function will have advanced the cursor until the 104 | // end of the frame. Since the cursor had position set to zero 105 | // before `Frame::check` was called, we obtain the length of the 106 | // frame by checking the cursor position. 107 | let len = buf.position() as usize; 108 | 109 | // Reset the position to zero before passing the cursor to 110 | // `Frame::parse`. 111 | buf.set_position(0); 112 | 113 | // Parse the frame from the buffer. This allocates the necessary 114 | // structures to represent the frame and returns the frame value. 115 | // 116 | // If the encoded frame representation is invalid, an error is 117 | // returned. This should terminate the **current** connection 118 | // but should not impact any other connected client. 119 | let frame = Frame::parse(&mut buf)?; 120 | 121 | // Discard the parsed data from the read buffer. 122 | // 123 | // When `advance` is called on the read buffer, all of the data 124 | // up to `len` is discarded. The details of how this works is 125 | // left to `BytesMut`. This is often done by moving an internal 126 | // cursor, but it may be done by reallocating and copying data. 127 | self.buffer.advance(len); 128 | 129 | // Return the parsed frame to the caller. 130 | Ok(Some(frame)) 131 | } 132 | // There is not enough data present in the read buffer to parse a 133 | // single frame. We must wait for more data to be received from the 134 | // socket. Reading from the socket will be done in the statement 135 | // after this `match`. 136 | // 137 | // We do not want to return `Err` from here as this "error" is an 138 | // expected runtime condition. 139 | Err(MiniRedisParseError::Incomplete) => Ok(None), 140 | // An error was encountered while parsing the frame. The connection 141 | // is now in an invalid state. Returning `Err` from here will result 142 | // in the connection being closed. 143 | Err(e) => Err(e.into()), 144 | } 145 | } 146 | 147 | /// Write a single `Frame` value to the underlying stream. 148 | /// 149 | /// The `Frame` value is written to the socket using the various `write_*` 150 | /// functions provided by `AsyncWrite`. Calling these functions directly on 151 | /// a `TcpStream` is **not** advised, as this will result in a large number of 152 | /// syscalls. However, it is fine to call these functions on a *buffered* 153 | /// write stream. The data will be written to the buffer. Once the buffer is 154 | /// full, it is flushed to the underlying socket. 155 | pub async fn write_frame(&mut self, frame: &Frame) -> Result<(), MiniRedisConnectionError> { 156 | // Arrays are encoded by encoding each entry. All other frame types are 157 | // considered literals. For now, mini-redis is not able to encode 158 | // recursive frame structures. See below for more details. 159 | match frame { 160 | Frame::Array(val) => { 161 | // Encode the frame type prefix. For an array, it is `*`. 162 | self.stream.write_u8(b'*').await?; 163 | 164 | // Encode the length of the array. 165 | self.write_decimal(val.len() as u64).await?; 166 | 167 | // Iterate and encode each entry in the array. 168 | // todo ? 169 | // for in &**val { 170 | for entry in val { 171 | self.write_value(entry).await?; 172 | } 173 | } 174 | // The frame type is a literal. Encode the value directly. 175 | _ => self.write_value(frame).await?, 176 | } 177 | 178 | // Ensure the encoded frame is written to the socket. The calls above 179 | // are to the buffered stream and writes. Calling `flush` writes the 180 | // remaining contents of the buffer to the socket. 181 | self.stream.flush().await.map_err(|e| e.into()) 182 | } 183 | 184 | /// Write a frame literal to the stream 185 | async fn write_value(&mut self, frame: &Frame) -> Result<(), MiniRedisConnectionError> { 186 | match frame { 187 | Frame::Simple(val) => { 188 | self.stream.write_u8(b'+').await?; 189 | self.stream.write_all(val.as_bytes()).await?; 190 | self.stream.write_all(b"\r\n").await?; 191 | } 192 | Frame::Error(val) => { 193 | self.stream.write_u8(b'-').await?; 194 | self.stream.write_all(val.as_bytes()).await?; 195 | self.stream.write_all(b"\r\n").await?; 196 | } 197 | Frame::Integer(val) => { 198 | self.stream.write_u8(b':').await?; 199 | self.write_decimal(*val).await?; 200 | } 201 | Frame::Null => { 202 | self.stream.write_all(b"$-1\r\n").await?; 203 | } 204 | Frame::Bulk(val) => { 205 | let len = val.len(); 206 | 207 | self.stream.write_u8(b'$').await?; 208 | self.write_decimal(len as u64).await?; 209 | self.stream.write_all(val).await?; 210 | self.stream.write_all(b"\r\n").await?; 211 | } 212 | // Encoding an `Array` from within a value cannot be done using a 213 | // recursive strategy. In general, async fns do not support 214 | // recursion. Mini-redis has not needed to encode nested arrays yet, 215 | // so for now it is skipped. 216 | Frame::Array(_val) => { 217 | warn!("unreachable code: recursive write_value: {:?}", _val); 218 | return Err(MiniRedisParseError::Unimplemented.into()); 219 | } 220 | } 221 | 222 | Ok(()) 223 | } 224 | 225 | /// Write a decimal frame to the stream 226 | async fn write_decimal(&mut self, val: u64) -> Result<(), MiniRedisConnectionError> { 227 | use std::io::Write; 228 | 229 | // Convert the value to a string 230 | let mut buf = [0u8; 20]; 231 | let mut buf = Cursor::new(&mut buf[..]); 232 | 233 | write!(&mut buf, "{}", val)?; 234 | 235 | let pos = buf.position() as usize; 236 | self.stream.write_all(&buf.get_ref()[..pos]).await?; 237 | self.stream.write_all(b"\r\n").await?; 238 | 239 | Ok(()) 240 | } 241 | } 242 | -------------------------------------------------------------------------------- /src/connection/frame.rs: -------------------------------------------------------------------------------- 1 | //! Provides a type representing a Redis protocol frame as well as utilities for 2 | //! parsing frames from a byte array. 3 | //! 4 | //! Redis serialization protocol (RESP) specification: 5 | //! https://redis.io/docs/reference/protocol-spec/ 6 | 7 | use std::convert::TryInto; 8 | use std::fmt; 9 | use std::io::Cursor; 10 | 11 | use bytes::{Buf, Bytes}; 12 | 13 | use crate::error::MiniRedisParseError; 14 | 15 | /// A frame in the Redis protocol. 16 | #[derive(Clone, Debug)] 17 | pub enum Frame { 18 | Simple(String), 19 | Error(String), 20 | Integer(u64), 21 | Bulk(Bytes), 22 | Null, 23 | Array(Vec), 24 | } 25 | 26 | impl PartialEq<&str> for Frame { 27 | fn eq(&self, other: &&str) -> bool { 28 | match self { 29 | Frame::Simple(s) => s.eq(other), 30 | Frame::Bulk(s) => s.eq(other), 31 | _ => false, 32 | } 33 | } 34 | } 35 | 36 | impl fmt::Display for Frame { 37 | fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { 38 | use std::str; 39 | 40 | match self { 41 | Frame::Simple(response) => response.fmt(fmt), 42 | Frame::Error(msg) => write!(fmt, "error: {}", msg), 43 | Frame::Integer(num) => num.fmt(fmt), 44 | Frame::Bulk(msg) => match str::from_utf8(msg) { 45 | Ok(string) => string.fmt(fmt), 46 | Err(_) => write!(fmt, "{:?}", msg), 47 | }, 48 | Frame::Null => "(nil)".fmt(fmt), 49 | Frame::Array(parts) => { 50 | for (i, part) in parts.iter().enumerate() { 51 | if i > 0 { 52 | write!(fmt, " ")?; 53 | part.fmt(fmt)?; 54 | } 55 | } 56 | 57 | Ok(()) 58 | } 59 | } 60 | } 61 | } 62 | 63 | impl Frame { 64 | /// Returns an empty array 65 | pub(crate) fn array() -> Frame { 66 | Frame::Array(vec![]) 67 | } 68 | 69 | /// Push a "bulk" frame into the array. `self` must be an Array frame. 70 | pub(crate) fn push_bulk(&mut self, bytes: Bytes) -> Result<(), MiniRedisParseError> { 71 | match self { 72 | Frame::Array(vec) => { 73 | vec.push(Frame::Bulk(bytes)); 74 | Ok(()) 75 | } 76 | _ => Err(MiniRedisParseError::ParseArrayFrame), 77 | } 78 | } 79 | 80 | /// Push an "integer" frame into the array. `self` must be an Array frame. 81 | pub(crate) fn push_int(&mut self, value: u64) -> Result<(), MiniRedisParseError> { 82 | match self { 83 | Frame::Array(vec) => { 84 | vec.push(Frame::Integer(value)); 85 | Ok(()) 86 | } 87 | _ => Err(MiniRedisParseError::ParseArrayFrame), 88 | } 89 | } 90 | 91 | /// Checks if an entire message can be decoded from `src` 92 | /// 93 | /// Redis serialization protocol (RESP) specification: 94 | /// https://redis.io/docs/reference/protocol-spec/ 95 | pub fn check(src: &mut Cursor<&[u8]>) -> Result<(), MiniRedisParseError> { 96 | match get_u8(src)? { 97 | b'+' => { 98 | get_line(src)?; 99 | Ok(()) 100 | } 101 | b'-' => { 102 | get_line(src)?; 103 | Ok(()) 104 | } 105 | b':' => { 106 | let _ = get_decimal(src)?; 107 | Ok(()) 108 | } 109 | b'$' => { 110 | if b'-' == peek_u8(src)? { 111 | // Skip '-1\r\n' 112 | skip(src, 4) 113 | } else { 114 | // Read the bulk string 115 | let len: usize = get_decimal(src)?.try_into()?; 116 | 117 | // skip that number of bytes + 2 (\r\n). 118 | skip(src, len + 2) 119 | } 120 | } 121 | b'*' => { 122 | let len = get_decimal(src)?; 123 | 124 | for _ in 0..len { 125 | Frame::check(src)?; 126 | } 127 | 128 | Ok(()) 129 | } 130 | actual => Err(MiniRedisParseError::Parse(format!( 131 | "protocol error; invalid frame type byte `{}`", 132 | actual 133 | ))), 134 | } 135 | } 136 | 137 | /// The message has already been validated with `check`, so parse the bytes to Frame 138 | /// 139 | /// Redis serialization protocol (RESP) specification: 140 | /// https://redis.io/docs/reference/protocol-spec/ 141 | pub fn parse(src: &mut Cursor<&[u8]>) -> Result { 142 | match get_u8(src)? { 143 | b'+' => { 144 | // Read the line and convert it to `Vec` 145 | let line = get_line(src)?.to_vec(); 146 | 147 | // Convert the line to a String 148 | let string = String::from_utf8(line)?; 149 | 150 | Ok(Frame::Simple(string)) 151 | } 152 | b'-' => { 153 | // Read the line and convert it to `Vec` 154 | let line = get_line(src)?.to_vec(); 155 | 156 | // Convert the line to a String 157 | let string = String::from_utf8(line)?; 158 | 159 | Ok(Frame::Error(string)) 160 | } 161 | b':' => { 162 | let len = get_decimal(src)?; 163 | Ok(Frame::Integer(len)) 164 | } 165 | b'$' => { 166 | if b'-' == peek_u8(src)? { 167 | let line = get_line(src)?; 168 | 169 | if line != b"-1" { 170 | return Err(MiniRedisParseError::Parse( 171 | "protocol error; invalid frame format".into(), 172 | )); 173 | } 174 | 175 | Ok(Frame::Null) 176 | } else { 177 | // Read the bulk string 178 | let len = get_decimal(src)?.try_into()?; 179 | let n = len + 2; 180 | 181 | if src.remaining() < n { 182 | return Err(MiniRedisParseError::Incomplete); 183 | } 184 | 185 | let data = Bytes::copy_from_slice(&src.chunk()[..len]); 186 | 187 | // skip that number of bytes + 2 (\r\n). 188 | skip(src, n)?; 189 | 190 | Ok(Frame::Bulk(data)) 191 | } 192 | } 193 | b'*' => { 194 | let len = get_decimal(src)?.try_into()?; 195 | let mut out = Vec::with_capacity(len); 196 | 197 | for _ in 0..len { 198 | out.push(Frame::parse(src)?); 199 | } 200 | 201 | Ok(Frame::Array(out)) 202 | } 203 | _ => Err(MiniRedisParseError::Unimplemented), 204 | } 205 | } 206 | } 207 | 208 | fn skip(src: &mut Cursor<&[u8]>, n: usize) -> Result<(), MiniRedisParseError> { 209 | if src.remaining() < n { 210 | return Err(MiniRedisParseError::Incomplete); 211 | } 212 | 213 | src.advance(n); 214 | Ok(()) 215 | } 216 | 217 | fn peek_u8(src: &mut Cursor<&[u8]>) -> Result { 218 | if !src.has_remaining() { 219 | return Err(MiniRedisParseError::Incomplete); 220 | } 221 | 222 | Ok(src.chunk()[0]) 223 | } 224 | 225 | fn get_u8(src: &mut Cursor<&[u8]>) -> Result { 226 | if !src.has_remaining() { 227 | return Err(MiniRedisParseError::Incomplete); 228 | } 229 | 230 | Ok(src.get_u8()) 231 | } 232 | 233 | /// Read a new-line terminated decimal 234 | fn get_decimal(src: &mut Cursor<&[u8]>) -> Result { 235 | use atoi::atoi; 236 | 237 | let line = get_line(src)?; 238 | 239 | atoi::(line).ok_or_else(|| { 240 | MiniRedisParseError::Parse("protocol error; invalid frame format to get decimal".into()) 241 | }) 242 | } 243 | 244 | /// Find a line in a frame 245 | fn get_line<'a>(src: &mut Cursor<&'a [u8]>) -> Result<&'a [u8], MiniRedisParseError> { 246 | // Scan the bytes directly 247 | let start = src.position() as usize; 248 | // Scan to the second to last byte 249 | let end = src.get_ref().len() - 1; 250 | 251 | for i in start..end { 252 | if src.get_ref()[i] == b'\r' && src.get_ref()[i + 1] == b'\n' { 253 | // We found a line, update the position to be *after* the \n 254 | src.set_position((i + 2) as u64); 255 | 256 | // Return the line 257 | return Ok(&src.get_ref()[start..i]); 258 | } 259 | } 260 | 261 | Err(MiniRedisParseError::Incomplete) 262 | } 263 | -------------------------------------------------------------------------------- /src/connection/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod connect; 2 | pub mod frame; 3 | pub mod parse; 4 | -------------------------------------------------------------------------------- /src/connection/parse.rs: -------------------------------------------------------------------------------- 1 | use std::vec; 2 | 3 | use bytes::Bytes; 4 | 5 | use crate::connection::frame::Frame; 6 | use crate::error::MiniRedisParseError; 7 | 8 | /// Utility for parsing a command 9 | /// 10 | /// Commands are represented as array frames. Each entry in the frame is a 11 | /// "token". A `Parse` is initialized with the array frame and provides a 12 | /// cursor-like API. Each command struct includes a `parse_frame` method that 13 | /// uses a `Parse` to extract its fields. 14 | #[derive(Debug)] 15 | pub(crate) struct Parse { 16 | /// Array frame iterator. 17 | parts: vec::IntoIter, 18 | } 19 | 20 | impl Parse { 21 | /// Create a new `Parse` to parse the contents of `frame`. 22 | /// 23 | /// Returns `Err` if `frame` is not an array frame. 24 | pub(crate) fn new(frame: Frame) -> Result { 25 | let array = match frame { 26 | Frame::Array(array) => array, 27 | frame => { 28 | return Err(MiniRedisParseError::Parse(format!( 29 | "protocol error; expected array, got {:?}", 30 | frame 31 | ))) 32 | } 33 | }; 34 | 35 | Ok(Parse { 36 | parts: array.into_iter(), 37 | }) 38 | } 39 | 40 | /// Return the next entry. Array frames are arrays of frames, so the next 41 | /// entry is a frame. 42 | fn next(&mut self) -> Result { 43 | self.parts.next().ok_or(MiniRedisParseError::EndOfStream) 44 | } 45 | 46 | /// Return the next entry as a string. 47 | /// 48 | /// If the next entry cannot be represented as a String, then an error is returned. 49 | pub(crate) fn next_string(&mut self) -> Result { 50 | match self.next()? { 51 | // Both `Simple` and `Bulk` representation may be strings. Strings 52 | // are parsed to UTF-8. 53 | // 54 | // While errors are stored as strings, they are considered separate 55 | // types. 56 | Frame::Simple(s) => Ok(s), 57 | Frame::Bulk(data) => std::str::from_utf8(&data[..]) 58 | .map(|s| s.to_string()) 59 | .map_err(|_| MiniRedisParseError::Parse("protocol error; invalid string".into())), 60 | frame => Err(MiniRedisParseError::Parse(format!( 61 | "protocol error; expected simple frame or bulk frame, got {:?}", 62 | frame 63 | ))), 64 | } 65 | } 66 | 67 | /// Return the next entry as raw bytes. 68 | /// 69 | /// If the next entry cannot be represented as raw bytes, an error is 70 | /// returned. 71 | pub(crate) fn next_bytes(&mut self) -> Result { 72 | match self.next()? { 73 | // Both `Simple` and `Bulk` representation may be raw bytes. 74 | // 75 | // Although errors are stored as strings and could be represented as 76 | // raw bytes, they are considered separate types. 77 | Frame::Simple(s) => Ok(Bytes::from(s.into_bytes())), 78 | Frame::Bulk(data) => Ok(data), 79 | frame => Err(MiniRedisParseError::Parse(format!( 80 | "protocol error; expected simple frame or bulk frame, got {:?}", 81 | frame 82 | ))), 83 | } 84 | } 85 | 86 | /// Return the next entry as an integer. 87 | /// 88 | /// This includes `Simple`, `Bulk`, and `Integer` frame types. `Simple` and 89 | /// `Bulk` frame types are parsed. 90 | /// 91 | /// If the next entry cannot be represented as an integer, then an error is 92 | /// returned. 93 | pub(crate) fn next_int(&mut self) -> Result { 94 | use atoi::atoi; 95 | 96 | match self.next()? { 97 | // An integer frame type is already stored as an integer. 98 | Frame::Integer(v) => Ok(v), 99 | // Simple and bulk frames must be parsed as integers. If the parsing 100 | // fails, an error is returned. 101 | Frame::Simple(data) => atoi::(data.as_bytes()) 102 | .ok_or_else(|| MiniRedisParseError::Parse("protocol error; invalid number".into())), 103 | Frame::Bulk(data) => atoi::(&data) 104 | .ok_or_else(|| MiniRedisParseError::Parse("protocol error; invalid number".into())), 105 | frame => Err(MiniRedisParseError::Parse(format!( 106 | "protocol error; expected int frame but got {:?}", 107 | frame 108 | ))), 109 | } 110 | } 111 | 112 | /// Ensure there are no more entries in the array 113 | pub(crate) fn finish(&mut self) -> Result<(), MiniRedisParseError> { 114 | if self.parts.next().is_none() { 115 | Ok(()) 116 | } else { 117 | Err(MiniRedisParseError::Parse( 118 | "protocol error; expected end of frame, but there was more".into(), 119 | )) 120 | } 121 | } 122 | } 123 | -------------------------------------------------------------------------------- /src/consts.rs: -------------------------------------------------------------------------------- 1 | /// Default port that a redis server listens on. 2 | /// 3 | /// Used if no port is specified. 4 | pub const DEFAULT_PORT: u16 = 6379; 5 | 6 | /// Maximum number of concurrent connections the redis server will accept. 7 | /// 8 | /// When this limit is reached, the server will stop accepting connections until 9 | /// an active connection terminates. 10 | pub const MAX_CONNECTIONS: usize = 1024; 11 | -------------------------------------------------------------------------------- /src/error.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | 3 | use thiserror::Error; 4 | 5 | #[derive(Error, Debug)] 6 | pub enum MiniRedisServerError { 7 | #[error(transparent)] 8 | IoError(#[from] io::Error), 9 | 10 | #[error(transparent)] 11 | Connect(#[from] MiniRedisConnectionError), 12 | 13 | #[error(transparent)] 14 | Parse(#[from] MiniRedisParseError), 15 | } 16 | 17 | #[derive(Error, Debug)] 18 | pub enum MiniRedisClientError { 19 | #[error(transparent)] 20 | Connect(#[from] MiniRedisConnectionError), 21 | 22 | #[error(transparent)] 23 | Parse(#[from] MiniRedisParseError), 24 | } 25 | 26 | /// Error encountered while parsing a frame. 27 | /// 28 | /// Only `EndOfStream` errors are handled at runtime. All other errors result in 29 | /// the connection being terminated. 30 | #[derive(Error, Debug)] 31 | pub enum MiniRedisParseError { 32 | #[error("invalid message encoding, parse failed")] 33 | Parse(String), 34 | 35 | /// Attempting to extract a value failed due to the frame being fully 36 | /// consumed. 37 | #[error("protocol error; unexpected end of stream")] 38 | EndOfStream, 39 | 40 | #[error("not enough data is available to parse a message")] 41 | Incomplete, 42 | 43 | #[error("unimplemented command")] 44 | Unimplemented, 45 | 46 | #[error("not an array frame")] 47 | ParseArrayFrame, 48 | 49 | #[error(transparent)] 50 | ParseInt(#[from] std::num::TryFromIntError), 51 | #[error(transparent)] 52 | ParseUtf8(#[from] std::string::FromUtf8Error), 53 | } 54 | 55 | #[derive(Error, Debug)] 56 | pub enum MiniRedisConnectionError { 57 | #[error("connection reset by peer")] 58 | Disconnect, 59 | 60 | #[error(transparent)] 61 | ParseFrame(#[from] MiniRedisParseError), 62 | 63 | #[error(transparent)] 64 | IoError(#[from] io::Error), 65 | 66 | #[error("command execute error")] 67 | CommandExecute(String), 68 | 69 | #[error("received next message failed, invalid frame type")] 70 | InvalidFrameType, 71 | 72 | #[error("invalid argument")] 73 | InvalidArgument(String), 74 | } 75 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | pub mod client; 2 | pub mod cmd; 3 | pub mod config; 4 | pub(crate) mod connection; 5 | pub mod consts; 6 | pub mod error; 7 | pub mod logger; 8 | pub mod server; 9 | mod storage; 10 | -------------------------------------------------------------------------------- /src/logger.rs: -------------------------------------------------------------------------------- 1 | use std::env; 2 | 3 | use log::{Level, LevelFilter, Metadata, Record}; 4 | 5 | use crate::config::LOG_LEVEL; 6 | 7 | struct Logger; 8 | 9 | pub fn init() { 10 | static LOGGER: Logger = Logger; 11 | log::set_logger(&LOGGER).unwrap(); 12 | 13 | let log_level: String = env::var(LOG_LEVEL).unwrap_or_else(|_| String::from("INFO")); 14 | log::set_max_level(match log_level.as_str() { 15 | "ERROR" => LevelFilter::Error, 16 | "WARN" => LevelFilter::Warn, 17 | "INFO" => LevelFilter::Info, 18 | "DEBUG" => LevelFilter::Debug, 19 | "TRACE" => LevelFilter::Trace, 20 | _ => LevelFilter::Info, 21 | }); 22 | } 23 | 24 | impl log::Log for Logger { 25 | fn enabled(&self, _metadata: &Metadata) -> bool { 26 | true 27 | } 28 | 29 | fn log(&self, record: &Record) { 30 | if !self.enabled(record.metadata()) { 31 | return; 32 | } 33 | 34 | let color = match record.level() { 35 | Level::Error => 31, // Red 36 | Level::Warn => 93, // BrightYellow 37 | Level::Info => 34, // Blue 38 | Level::Debug => 32, // Green 39 | Level::Trace => 90, // BrightBlack 40 | }; 41 | 42 | println!( 43 | "\u{1B}[{}m[{:>5}]: {} - {}\u{1B}[0m", 44 | color, 45 | record.level(), 46 | record.target(), 47 | record.args(), 48 | ); 49 | } 50 | 51 | fn flush(&self) {} 52 | } 53 | -------------------------------------------------------------------------------- /src/server/handler.rs: -------------------------------------------------------------------------------- 1 | use log::debug; 2 | use tokio::sync::mpsc; 3 | 4 | use crate::cmd::Command; 5 | use crate::connection::connect::Connection; 6 | use crate::error::MiniRedisConnectionError; 7 | use crate::server::shutdown::Shutdown; 8 | use crate::storage::db::Db; 9 | 10 | /// Per-connection handler. Reads requests from `connection` and applies the 11 | /// commands to `db`. 12 | #[derive(Debug)] 13 | pub(crate) struct Handler { 14 | /// Shared database handle. 15 | /// 16 | /// When a command is received from `connection`, it is applied with `db`. 17 | /// The implementation of the command is in the `cmd` module. Each command 18 | /// will need to interact with `db` in order to complete the work. 19 | pub(crate) db: Db, 20 | 21 | /// The TCP connection decorated with the redis protocol encoder / decoder 22 | /// implemented using a buffered `TcpStream`. 23 | /// 24 | /// When `Listener` receives an inbound connection, the `TcpStream` is 25 | /// passed to `Connection::new`, which initializes the associated buffers. 26 | /// `Connection` allows the handler to operate at the "frame" level and keep 27 | /// the byte level protocol parsing details encapsulated in `Connection`. 28 | pub(crate) connection: Connection, 29 | 30 | /// Listen for shutdown notifications. 31 | /// 32 | /// A wrapper around the `broadcast::Receiver` paired with the sender in 33 | /// `Listener`. The connection handler processes requests from the 34 | /// connection until the peer disconnects **or** a shutdown notification is 35 | /// received from `shutdown`. In the latter case, any in-flight work being 36 | /// processed for the peer is continued until it reaches a safe state, at 37 | /// which point the connection is terminated. 38 | pub(crate) shutdown: Shutdown, 39 | 40 | /// Not used directly. Instead, when `Handler` is dropped...? 41 | pub(crate) _shutdown_complete: mpsc::Sender<()>, 42 | } 43 | 44 | impl Handler { 45 | /// Process a single connection. 46 | /// 47 | /// Request frames are read from the socket and processed. Responses are 48 | /// written back to the socket. 49 | /// 50 | /// Currently, pipelining is not implemented. Pipelining is the ability to 51 | /// process more than one request concurrently per connection without 52 | /// interleaving frames. See for more details: 53 | /// https://redis.io/topics/pipelining 54 | /// 55 | /// When the shutdown signal is received, the connection is processed until 56 | /// it reaches a safe state, at which point it is terminated. 57 | pub(crate) async fn run(&mut self) -> Result<(), MiniRedisConnectionError> { 58 | // As long as the shutdown signal has not been received, try to read a 59 | // new request frame. 60 | while !self.shutdown.is_shutdown() { 61 | // While reading a request frame, also listen for the shutdown 62 | // signal. 63 | let maybe_frame = tokio::select! { 64 | res = self.connection.read_frame() => res?, 65 | _ = self.shutdown.recv() => { 66 | // If a shutdown signal is received, return from `run`. 67 | // This will result in the task terminating. 68 | return Ok(()); 69 | } 70 | }; 71 | 72 | // If `None` is returned from `read_frame()` then the peer closed 73 | // the socket. There is no further work to do and the task can be 74 | // terminated. 75 | let frame = match maybe_frame { 76 | Some(frame) => frame, 77 | None => { 78 | debug!("peer closed the socket, return"); 79 | return Ok(()); 80 | } 81 | }; 82 | 83 | // Convert the redis frame into a command struct. This returns an 84 | // error if the frame is not a valid redis command or it is an 85 | // unsupported command. 86 | let cmd = Command::from_frame(frame)?; 87 | 88 | // Logs the `cmd` object. 89 | debug!("received command: {:?}", cmd); 90 | 91 | // Perform the work needed to apply the command. This may mutate the 92 | // database state as a result. 93 | // 94 | // The connection is passed into the apply function which allows the 95 | // command to write response frames directly to the connection. In 96 | // the case of pub/sub, multiple frames may be send back to the 97 | // peer. 98 | cmd.apply(&self.db, &mut self.connection, &mut self.shutdown) 99 | .await?; 100 | } 101 | 102 | Ok(()) 103 | } 104 | } 105 | -------------------------------------------------------------------------------- /src/server/listener.rs: -------------------------------------------------------------------------------- 1 | use std::sync::Arc; 2 | use std::time::Duration; 3 | 4 | use log::{error, info}; 5 | use tokio::net::{TcpListener, TcpStream}; 6 | use tokio::sync::{broadcast, mpsc, Semaphore}; 7 | use tokio::time; 8 | 9 | use crate::connection::connect::Connection; 10 | use crate::error::MiniRedisConnectionError; 11 | use crate::server::handler::Handler; 12 | use crate::server::shutdown::Shutdown; 13 | use crate::storage::db::DbDropGuard; 14 | 15 | /// Server listener state. Created in the `run` call. It includes a `run` method 16 | /// which performs the TCP listening and initialization of per-connection state. 17 | #[derive(Debug)] 18 | pub(crate) struct Listener { 19 | /// Shared database handle. 20 | /// 21 | /// Contains the key / value store as well as the broadcast channels for 22 | /// pub/sub. 23 | /// 24 | /// This holds a wrapper around an `Arc`. The internal `Db` can be 25 | /// retrieved and passed into the per connection state (`Handler`). 26 | pub(crate) db_holder: DbDropGuard, 27 | 28 | /// TCP listener supplied by the `run` caller. 29 | pub(crate) listener: TcpListener, 30 | 31 | /// Limit the max number of connections. 32 | /// 33 | /// A `Semaphore` is used to limit the max number of connections. Before 34 | /// attempting to accept a new connection, a permit is acquired from the 35 | /// semaphore. If none are available, the listener waits for one. 36 | /// 37 | /// When handlers complete processing a connection, the permit is returned 38 | /// to the semaphore. 39 | pub(crate) limit_connections: Arc, 40 | 41 | /// Broadcasts a shutdown signal to all active connections. 42 | /// 43 | /// The initial `shutdown` trigger is provided by the `run` caller. The 44 | /// server is responsible for gracefully shutting down active connections. 45 | /// When a connection task is spawned, it is passed a broadcast receiver 46 | /// handle. When a graceful shutdown is initiated, a `()` value is sent via 47 | /// the broadcast::Sender. Each active connection receives it, reaches a 48 | /// safe terminal state, and completes the task. 49 | pub(crate) notify_shutdown: broadcast::Sender<()>, 50 | 51 | /// Used as part of the graceful shutdown process to wait for client 52 | /// connections to complete processing. 53 | /// 54 | /// Tokio channels are closed once all `Sender` handles go out of scope. 55 | /// When a channel is closed, the receiver receives `None`. This is 56 | /// leveraged to detect all connection handlers completing. When a 57 | /// connection handler is initialized, it is assigned a clone of 58 | /// `shutdown_complete_tx`. When the listener shuts down, it drops the 59 | /// sender held by this `shutdown_complete_tx` field. Once all handler tasks 60 | /// complete, all clones of the `Sender` are also dropped. This results in 61 | /// `shutdown_complete_rx.recv()` completing with `None`. At this point, it 62 | /// is safe to exit the server process. 63 | pub(crate) shutdown_complete_rx: mpsc::Receiver<()>, 64 | pub(crate) shutdown_complete_tx: mpsc::Sender<()>, 65 | } 66 | 67 | impl Listener { 68 | /// Run the server 69 | /// 70 | /// Listen for inbound connections. For each inbound connection, spawn a 71 | /// task to process that connection. 72 | /// 73 | /// # Errors 74 | /// 75 | /// Returns `Err` if accepting returns an error. This can happen for a 76 | /// number reasons that resolve over time. For example, if the underlying 77 | /// operating system has reached an internal limit for max number of 78 | /// sockets, accept will fail. 79 | /// 80 | /// The process is not able to detect when a transient error resolves 81 | /// itself. One strategy for handling this is to implement a back off 82 | /// strategy, which is what we do here. 83 | pub(crate) async fn run(&mut self) -> Result<(), MiniRedisConnectionError> { 84 | info!("server started, accepting inbound connections"); 85 | 86 | loop { 87 | // Wait for a permit to become available 88 | // 89 | // `acquire_owned` returns a permit that is bound to the semaphore. 90 | // When the permit value is dropped, it is automatically returned 91 | // to the semaphore. 92 | // 93 | // `acquire_owned()` returns `Err` when the semaphore has been 94 | // closed. We don't ever close the semaphore, so `unwrap()` is safe. 95 | let permit = self 96 | .limit_connections 97 | .clone() 98 | .acquire_owned() 99 | .await 100 | .unwrap(); 101 | 102 | // Accept a new socket. This will attempt to perform error handling. 103 | // The `accept` method internally attempts to recover errors, so an 104 | // error here is non-recoverable. 105 | let socket = self.accept().await?; 106 | 107 | // Create the necessary per-connection handler state. 108 | let mut handler = Handler { 109 | // Get a handle to the shared database. 110 | db: self.db_holder.db(), 111 | 112 | // Initialize the connection state. This allocates read/write 113 | // buffers to perform redis protocol frame parsing. 114 | connection: Connection::new(socket), 115 | 116 | // Receive shutdown notifications. 117 | shutdown: Shutdown::new(self.notify_shutdown.subscribe()), 118 | 119 | // Notifies the receiver half once all clones are dropped. 120 | _shutdown_complete: self.shutdown_complete_tx.clone(), 121 | }; 122 | 123 | // Spawn a new task to process the connections. Tokio tasks are like 124 | // asynchronous green threads and are executed concurrently. 125 | tokio::spawn(async move { 126 | // Process the connection. If an error is encountered, log it. 127 | if let Err(err) = handler.run().await { 128 | error!("connection error：{:?}", err); 129 | } 130 | // Move the permit into the task and drop it after completion. 131 | // This returns the permit back to the semaphore. 132 | drop(permit); 133 | }); 134 | } 135 | } 136 | 137 | /// Accept an inbound connection. 138 | /// 139 | /// Errors are handled by backing off and retrying. An exponential backoff 140 | /// strategy is used. After the first failure, the task waits for 1 second. 141 | /// After the second failure, the task waits for 2 seconds. Each subsequent 142 | /// failure doubles the wait time. If accepting fails on the 6th try after 143 | /// waiting for 64 seconds, then this function returns with an error. 144 | async fn accept(&mut self) -> Result { 145 | let mut backoff = 1; 146 | 147 | // Try to accept a few times 148 | loop { 149 | // Perform the accept operation. If a socket is successfully 150 | // accepted, return it. Otherwise, save the error. 151 | match self.listener.accept().await { 152 | Ok((socket, _)) => return Ok(socket), 153 | Err(err) => { 154 | if backoff > 64 { 155 | // Accept has failed too many times. Return the error. 156 | error!("failed to accept socket after retry: {}", err); 157 | return Err(err.into()); 158 | } else { 159 | error!("failed to accept socket: {}", err); 160 | } 161 | } 162 | } 163 | 164 | // Pause execution until the back off period elapses. 165 | time::sleep(Duration::from_secs(backoff)).await; 166 | 167 | // Double the back off 168 | backoff <<= 2; 169 | } 170 | } 171 | } 172 | -------------------------------------------------------------------------------- /src/server/mod.rs: -------------------------------------------------------------------------------- 1 | //! Minimal Redis server implementation 2 | //! 3 | //! Provides an async `run` function that listens for inbound connections, 4 | //! spawning one task per connection. 5 | 6 | use std::future::Future; 7 | use std::sync::Arc; 8 | 9 | use log::{debug, error, info}; 10 | use tokio::net::TcpListener; 11 | use tokio::sync::{broadcast, mpsc, Semaphore}; 12 | 13 | use crate::consts::MAX_CONNECTIONS; 14 | use crate::server::listener::Listener; 15 | use crate::storage::db::DbDropGuard; 16 | 17 | mod handler; 18 | pub(crate) mod listener; 19 | pub(crate) mod shutdown; 20 | 21 | /// Run the mini-redis server. 22 | /// 23 | /// Accepts connections from the supplied listener. For each inbound connection, 24 | /// a task is spawned to handle that connection. The server runs until the 25 | /// `shutdown` future completes, at which point the server shuts down 26 | /// gracefully. 27 | /// 28 | /// `tokio::signal::ctrl_c()` can be used as the `shutdown` argument. This will 29 | /// listen for a SIGINT signal. 30 | pub async fn run(listener: TcpListener, shutdown: impl Future) { 31 | info!( 32 | "mini-redis server started listen on: {}", 33 | listener.local_addr().unwrap() 34 | ); 35 | 36 | // When the provided `shutdown` future completes, we must send a shutdown 37 | // message to all active connections. We use a broadcast channel for this 38 | // purpose. The call below ignores the receiver of the broadcast pair, and when 39 | // a receiver is needed, the subscribe() method on the sender is used to create one. 40 | let (notify_shutdown, _) = broadcast::channel(1); 41 | let (shutdown_complete_tx, shutdown_complete_rx) = mpsc::channel(1); 42 | 43 | // Initialize the listener state 44 | let mut server = Listener { 45 | listener, 46 | db_holder: DbDropGuard::new(), 47 | limit_connections: Arc::new(Semaphore::new(MAX_CONNECTIONS)), 48 | notify_shutdown, 49 | shutdown_complete_tx, 50 | shutdown_complete_rx, 51 | }; 52 | 53 | // Concurrently run the server and listen for the `shutdown` signal. The 54 | // server task runs until an error is encountered, so under normal 55 | // circumstances, this `select!` statement runs until the `shutdown` signal 56 | // is received. 57 | // 58 | // `select!` statements are written in the form of: 59 | // 60 | // ```text 61 | // = => 62 | // ``` 63 | // 64 | // All `` statements are executed concurrently. Once the **first** 65 | // op completes, its associated `` is 66 | // performed. 67 | // 68 | // The `select!` macro is a foundational building block for writing 69 | // asynchronous Rust. See the API docs for more details: 70 | // 71 | // [select](https://docs.rs/tokio/*/tokio/macro.select.html) 72 | tokio::select! { 73 | res = server.run() => { 74 | // If an error is received here, accepting connections from the TCP 75 | // listener failed multiple times and the server is giving up. 76 | // 77 | // Errors encountered when handling individual connections do not 78 | // bubble up to this point. 79 | if let Err(err) = res { 80 | error!("failed to accept: {:?}", err); 81 | } 82 | } 83 | _ = shutdown => { 84 | // The shutdown signal has been received 85 | debug!("server is about to shutdown"); 86 | } 87 | } 88 | 89 | // Extract the `shutdown_complete` receiver and transmitter 90 | // explicitly drop `shutdown_transmitter`. This is important, as the 91 | // `.await` below would otherwise never complete. 92 | let Listener { 93 | mut shutdown_complete_rx, 94 | shutdown_complete_tx, 95 | notify_shutdown, 96 | .. 97 | } = server; 98 | 99 | // When `notify_shutdown` is dropped, all tasks which have `subscribe`d will 100 | // receive the shutdown signal and can exit 101 | drop(notify_shutdown); 102 | // Drop final `Sender` so the `Receiver` below can complete 103 | drop(shutdown_complete_tx); 104 | 105 | // Wait for all active connections to finish processing. As the `Sender` 106 | // handle held by the listener has been dropped above, the only remaining 107 | // `Sender` instances are held by connection handler tasks. When those drop, 108 | // the `mpsc` channel will close and `recv()` will return `None`. 109 | let _ = shutdown_complete_rx.recv().await; 110 | } 111 | -------------------------------------------------------------------------------- /src/server/shutdown.rs: -------------------------------------------------------------------------------- 1 | use tokio::sync::broadcast; 2 | 3 | /// Listens for the server shutdown signal. 4 | /// 5 | /// Shutdown is signalled using a `broadcast::Receiver`. Only a single value is 6 | /// ever sent. Once a value has been sent via the broadcast channel, the server 7 | /// should shutdown. 8 | /// 9 | /// The `Shutdown` struct listens for the signal and tracks that the signal has 10 | /// been received. Callers may query for whether the shutdown signal has been 11 | /// received or not. 12 | #[derive(Debug)] 13 | pub(crate) struct Shutdown { 14 | /// `true` if the shutdown signal has been received 15 | shutdown: bool, 16 | 17 | /// The receive half of the channel used to listen for shutdown. 18 | notify: broadcast::Receiver<()>, 19 | } 20 | 21 | impl Shutdown { 22 | /// Create a new `Shutdown` backed by the given `broadcast::Receiver`. 23 | pub(crate) fn new(notify: broadcast::Receiver<()>) -> Shutdown { 24 | Shutdown { 25 | shutdown: false, 26 | notify, 27 | } 28 | } 29 | 30 | /// Returns `true` if the shutdown signal has been received. 31 | pub(crate) fn is_shutdown(&self) -> bool { 32 | self.shutdown 33 | } 34 | 35 | /// Receive the shutdown notice, waiting if necessary. 36 | pub(crate) async fn recv(&mut self) { 37 | // If the shutdown signal has already been received, then return immediately. 38 | if self.shutdown { 39 | return; 40 | } 41 | 42 | // Cannot receive a "lag error" as only one value is ever sent. 43 | let _ = self.notify.recv().await; 44 | 45 | // Remember that the signal has been received. 46 | self.shutdown = true; 47 | } 48 | } 49 | -------------------------------------------------------------------------------- /src/storage/db.rs: -------------------------------------------------------------------------------- 1 | use std::sync::{Arc, Mutex}; 2 | 3 | use bytes::Bytes; 4 | use log::{debug, info}; 5 | use tokio::sync::{broadcast, Notify}; 6 | use tokio::time::{self, Duration, Instant}; 7 | 8 | use crate::storage::store::{Entry, Store}; 9 | use crate::storage::traits::KvStore; 10 | 11 | /// A wrapper around a `Db` instance. This exists to allow orderly cleanup 12 | /// of the `Db` by signalling the background purge task to shut down when 13 | /// this struct is dropped. 14 | #[derive(Debug)] 15 | pub(crate) struct DbDropGuard { 16 | /// The `Db` instance that will be shut down when this `DbHolder` struct 17 | /// is dropped. 18 | db: Db, 19 | } 20 | 21 | impl DbDropGuard { 22 | /// Create a new `DbHolder`, wrapping a `Db` instance. When this is dropped 23 | /// the `Db`'s purge task will be shut down. 24 | pub(crate) fn new() -> DbDropGuard { 25 | DbDropGuard { db: Db::new() } 26 | } 27 | 28 | /// Get the shared database. Internally, this is an `Arc`, 29 | /// so a clone only increments the ref count. 30 | pub(crate) fn db(&self) -> Db { 31 | self.db.clone() 32 | } 33 | } 34 | 35 | impl Drop for DbDropGuard { 36 | fn drop(&mut self) { 37 | // Signal the 'Db' instance to shut down the task that purges expired keys 38 | self.db.shutdown_purge_task(); 39 | } 40 | } 41 | 42 | /// Server store shared across all connections. 43 | /// 44 | /// `Db` contains a `HashMap` storing the key/value data and all 45 | /// `broadcast::Sender` values for active pub/sub channels. 46 | /// 47 | /// A `Db` instance is a handle to shared store. Cloning `Db` is shallow and 48 | /// only incurs an atomic ref count increment. 49 | /// 50 | /// When a `Db` value is created, a background task is spawned. This task is 51 | /// used to expire values after the requested duration has elapsed. The task 52 | /// runs until all instances of `Db` are dropped, at which point the task 53 | /// terminates. 54 | #[derive(Debug, Clone)] 55 | pub(crate) struct Db { 56 | /// Handle to shared store. The background task will also have an 57 | /// `Arc`. 58 | shared: Arc, 59 | } 60 | 61 | impl Db { 62 | /// Create a new, empty, `Db` instance. Allocates shared store and spawns a 63 | /// background task to manage key expiration. 64 | pub(crate) fn new() -> Db { 65 | let shared = Arc::new(SharedDb::new()); 66 | 67 | // Start the background task. 68 | tokio::spawn(Db::purge_expired_tasks(shared.clone())); 69 | 70 | Db { shared } 71 | } 72 | 73 | /// Routine executed by the background task. 74 | /// 75 | /// Wait to be notified. On notification, purge any expired keys from the shared 76 | /// store handle. If `shutdown` is set, terminate the task. 77 | async fn purge_expired_tasks(shared: Arc) { 78 | // If the shutdown flag is set, then the task should exit. 79 | while !shared.is_shutdown() { 80 | // Purge all keys that are expired. The function returns the instant at 81 | // which the **next** key will expire. The worker should wait until the 82 | // instant has passed then purge again. 83 | if let Some(when) = shared.purge_expired_keys() { 84 | // Wait until the next key expires **or** until the background task 85 | // is notified. If the task is notified, then it must reload its 86 | // store as new keys have been set to expire early. This is done by 87 | // looping. 88 | tokio::select! { 89 | _ = time::sleep_until(when) => {} 90 | _ = shared.background_task.notified() => {} 91 | } 92 | } else { 93 | // There are no keys expiring in the future. Wait until the task is 94 | // notified. 95 | shared.background_task.notified().await; 96 | } 97 | } 98 | 99 | info!("Purge background task shut down") 100 | } 101 | 102 | /// Signals the purge background task to shut down. This is called by the 103 | /// `DbShutdown`s `Drop` implementation. 104 | fn shutdown_purge_task(&self) { 105 | // The background task must be signaled to shut down. This is done by 106 | // setting `Store::shutdown` to `true` and signalling the task. 107 | let mut store = self.shared.store.lock().unwrap(); 108 | store.shutdown = true; 109 | 110 | // Drop the lock before signalling the background task. This helps 111 | // reduce lock contention by ensuring the background task doesn't 112 | // wake up only to be unable to acquire the mutex. 113 | drop(store); 114 | self.shared.background_task.notify_one(); 115 | } 116 | } 117 | 118 | impl KvStore for Db { 119 | /// Get the value associated with a key. 120 | /// 121 | /// Returns `None` if there is no value associated with the key. This may be 122 | /// due to never having assigned a value to the key or a previously assigned 123 | /// value expired. 124 | fn get(&self, key: &str) -> Option { 125 | // Acquire the lock, get the entry and clone the value. 126 | // 127 | // Because data is stored using `Bytes`, a clone here is a shallow 128 | // clone. Data is not copied. 129 | let store = self.shared.store.lock().unwrap(); 130 | store.entries.get(key).map(|entry| entry.data.clone()) 131 | } 132 | 133 | /// Set the value associated with a key along with an optional expiration 134 | /// Duration. 135 | /// 136 | /// If a value is already associated with the key, it is removed. 137 | fn set(&self, key: String, value: Bytes, expire: Option) { 138 | let mut store = self.shared.store.lock().unwrap(); 139 | 140 | // Get and increment the next insertion ID. Guarded by the lock, this 141 | // ensures a unique identifier is associated with each `set` operation. 142 | let id = store.next_id; 143 | store.next_id += 1; 144 | 145 | // If this `set` becomes the key that expires **next**, the background 146 | // task needs to be notified so it can update its state. 147 | // 148 | // Whether or not the task needs to be notified is computed during the 149 | // `set` routine. 150 | let mut notify = false; 151 | 152 | let expires_at = expire.map(|duration| { 153 | // `Instant` at which the key expires. 154 | let when = Instant::now() + duration; 155 | 156 | // Only notify the worker task if the newly inserted expiration is the 157 | // **next** key to evict. In this case, the worker needs to be woken up 158 | // to update its state. 159 | notify = store 160 | .next_expiration() 161 | .map(|expiration| expiration > when) 162 | .unwrap_or(true); 163 | 164 | // Track the expiration. 165 | store.expirations.insert((when, id), key.clone()); 166 | when 167 | }); 168 | 169 | // Insert the entry into the `HashMap`. 170 | let prev = store.entries.insert( 171 | key, 172 | Entry { 173 | id, 174 | data: value, 175 | expires_at, 176 | }, 177 | ); 178 | 179 | // If there was a value previously associated with the key **and** it 180 | // had an expiration time. The associated entry in the `expirations` map 181 | // must also be removed. This avoids leaking data. 182 | if let Some(prev) = prev { 183 | if let Some(when) = prev.expires_at { 184 | // clear expiration 185 | store.expirations.remove(&(when, prev.id)); 186 | } 187 | } 188 | 189 | // Release the mutex before notifying the background task. This helps 190 | // reduce contention by avoiding the background task waking up only to 191 | // be unable to acquire the mutex due to this function still holding it. 192 | drop(store); 193 | 194 | if notify { 195 | // Finally, only notify the background task if it needs to update 196 | // its state to reflect a new expiration. 197 | self.shared.background_task.notify_one(); 198 | } 199 | } 200 | 201 | /// Returns a `Receiver` for the requested channel. 202 | /// 203 | /// The returned `Receiver` is used to receive values broadcast by `PUBLISH` 204 | /// commands. 205 | fn subscribe(&self, key: String) -> broadcast::Receiver { 206 | use std::collections::hash_map::Entry; 207 | 208 | // Acquire the mutex 209 | let mut store = self.shared.store.lock().unwrap(); 210 | 211 | // If there is no entry for the requested channel, then create a new 212 | // broadcast channel and associate it with the key. If one already 213 | // exists, return an associated receiver. 214 | match store.pub_sub.entry(key) { 215 | Entry::Occupied(e) => e.get().subscribe(), 216 | Entry::Vacant(e) => { 217 | // No broadcast channel exists yet, so create one. 218 | // 219 | // The channel is created with a capacity of `1024` messages. A 220 | // message is stored in the channel until **all** subscribers 221 | // have seen it. This means that a slow subscriber could result 222 | // in messages being held indefinitely. 223 | // 224 | // When the channel's capacity fills up, publishing will result 225 | // in old messages being dropped. This prevents slow consumers 226 | // from blocking the entire system. 227 | let (tx, rx) = broadcast::channel(1024); 228 | e.insert(tx); 229 | rx 230 | } 231 | } 232 | } 233 | 234 | /// Publish a message to the channel. Returns the number of subscribers 235 | /// listening on the channel. 236 | fn publish(&self, key: &str, value: Bytes) -> usize { 237 | debug!("publish: (key={}, len(value)={})", key, value.len()); 238 | 239 | let state = self.shared.store.lock().unwrap(); 240 | 241 | state 242 | .pub_sub 243 | .get(key) 244 | // On a successful message send on the broadcast channel, the number 245 | // of subscribers is returned. An error indicates there are no 246 | // receivers, in which case, `0` should be returned. 247 | .map(|tx| tx.send(value).unwrap_or(0)) 248 | // If there is no entry for the channel key, then there are no 249 | // subscribers. In this case, return `0`. 250 | .unwrap_or(0) 251 | } 252 | } 253 | 254 | #[derive(Debug)] 255 | struct SharedDb { 256 | /// The shared store is guarded by a mutex. This is a `std::sync::Mutex` and 257 | /// not a Tokio mutex. This is because there are no asynchronous operations 258 | /// being performed while holding the mutex. Additionally, the critical 259 | /// sections are very small. 260 | /// 261 | /// A Tokio mutex is mostly intended to be used when locks need to be held 262 | /// across `.await` yield points. All other cases are **usually** best 263 | /// served by a std mutex. If the critical section does not include any 264 | /// async operations but is long (CPU intensive or performing blocking 265 | /// operations), then the entire operation, including waiting for the mutex, 266 | /// is considered a "blocking" operation and `tokio::task::spawn_blocking` 267 | /// should be used. 268 | store: Mutex, 269 | 270 | /// Notifies the background task handling entry expiration. The background 271 | /// task waits on this to be notified, then checks for expired values or the 272 | /// shutdown signal. 273 | background_task: Notify, 274 | } 275 | 276 | impl SharedDb { 277 | fn new() -> Self { 278 | SharedDb { 279 | store: Mutex::new(Store::new()), 280 | background_task: Notify::new(), 281 | } 282 | } 283 | 284 | /// Purge all expired keys and return the `Instant` at which the **next** 285 | /// key will expire. The background task will sleep until this instant. 286 | fn purge_expired_keys(&self) -> Option { 287 | let mut store = self.store.lock().unwrap(); 288 | 289 | if store.shutdown { 290 | // The database is shutting down. All handles to the shared store 291 | // have dropped. The background task should exit. 292 | return None; 293 | } 294 | 295 | // This is needed to make the borrow checker happy. In short, `lock()` 296 | // returns a `MutexGuard` and not a `&mut Store`. The borrow checker is 297 | // not able to see "through" the mutex guard and determine that it is 298 | // safe to access both `store.expirations` and `store.entries` mutably, 299 | // so we get a "real" mutable reference to `Store` outside of the loop. 300 | let store = &mut *store; 301 | 302 | // Find all keys scheduled to expire **before** now. 303 | let now = Instant::now(); 304 | while let Some((&(when, id), key)) = store.expirations.iter().next() { 305 | if when > now { 306 | // Done purging, `when` is the instant at which the next key 307 | // expires. The worker task will wait until this instant. 308 | return Some(when); 309 | } 310 | 311 | // The key expired, remove it 312 | store.entries.remove(key); 313 | store.expirations.remove(&(when, id)); 314 | } 315 | 316 | None 317 | } 318 | 319 | /// Returns `true` if the database is shutting down 320 | /// 321 | /// The `shutdown` flag is set when all `Db` values have dropped, indicating 322 | /// that the shared store can no longer be accessed. 323 | fn is_shutdown(&self) -> bool { 324 | self.store.lock().unwrap().shutdown 325 | } 326 | } 327 | -------------------------------------------------------------------------------- /src/storage/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod db; 2 | pub mod store; 3 | pub(crate) mod traits; 4 | -------------------------------------------------------------------------------- /src/storage/store.rs: -------------------------------------------------------------------------------- 1 | //! Core storage implementation for mini-redis 2 | use std::collections::{BTreeMap, HashMap}; 3 | 4 | use bytes::Bytes; 5 | use tokio::sync::broadcast; 6 | use tokio::time::Instant; 7 | 8 | #[derive(Debug)] 9 | pub(crate) struct Store { 10 | /// The key-value data. We are not trying to do anything fancy so a 11 | /// `std::collections::HashMap` works fine. 12 | /// For production implementation, more complex structure can be used! 13 | pub(crate) entries: HashMap, 14 | 15 | /// The pub/sub key-space. Redis uses a **separate** key space for key-value 16 | /// and pub/sub. `mini-redis` handles this by using a separate `HashMap`. 17 | pub(crate) pub_sub: HashMap>, 18 | 19 | /// Tracks key TTLs. 20 | /// 21 | /// A `BTreeMap` is used to maintain expirations sorted by when they expire. 22 | /// This allows the background task to iterate this map to find the value 23 | /// expiring next. 24 | /// 25 | /// While highly unlikely, it is possible for more than one expiration to be 26 | /// created for the same instant. Because of this, the `Instant` is 27 | /// insufficient for the key. A unique expiration identifier (`u64`) is used 28 | /// to break these ties. 29 | pub(crate) expirations: BTreeMap<(Instant, u64), String>, 30 | 31 | /// Identifier to use for the next expiration. Each expiration is associated 32 | /// with a unique identifier. See above for why. 33 | pub(crate) next_id: u64, 34 | 35 | /// True when the Db instance is shutting down. This happens when all `Db` 36 | /// values drop. Setting this to `true` signals to the background task to 37 | /// exit. 38 | pub(crate) shutdown: bool, 39 | } 40 | 41 | /// Entry in the key-value store 42 | #[derive(Debug)] 43 | pub(crate) struct Entry { 44 | /// Uniquely identifies this entry. 45 | pub(crate) id: u64, 46 | 47 | /// Stored data 48 | pub(crate) data: Bytes, 49 | 50 | /// Instant at which the entry expires and should be removed from the 51 | /// database. 52 | pub(crate) expires_at: Option, 53 | } 54 | 55 | impl Store { 56 | pub(crate) fn new() -> Store { 57 | Store { 58 | entries: HashMap::new(), 59 | pub_sub: HashMap::new(), 60 | expirations: BTreeMap::new(), 61 | next_id: 0, 62 | shutdown: false, 63 | } 64 | } 65 | 66 | /// Get the next expiration instant for notify 67 | pub(crate) fn next_expiration(&self) -> Option { 68 | self.expirations.keys().next().map(|expire| expire.0) 69 | } 70 | } 71 | -------------------------------------------------------------------------------- /src/storage/traits.rs: -------------------------------------------------------------------------------- 1 | use bytes::Bytes; 2 | use tokio::sync::broadcast; 3 | use tokio::time::Duration; 4 | 5 | pub trait KvStore { 6 | fn get(&self, key: &str) -> Option; 7 | 8 | /// Set the value associated with a key along with an optional expiration 9 | /// Duration. 10 | /// 11 | /// If a value is already associated with the key, it is removed. 12 | fn set(&self, key: String, value: Bytes, expire: Option); 13 | 14 | /// Returns a `Receiver` for the requested channel. 15 | /// 16 | /// The returned `Receiver` is used to receive values broadcast by `PUBLISH` 17 | /// commands. 18 | fn subscribe(&self, key: String) -> broadcast::Receiver; 19 | 20 | /// Publish a message to the channel. Returns the number of subscribers 21 | /// listening on the channel. 22 | fn publish(&self, key: &str, value: Bytes) -> usize; 23 | } 24 | -------------------------------------------------------------------------------- /tests/client.rs: -------------------------------------------------------------------------------- 1 | use mini_redis::{client, server}; 2 | use std::net::SocketAddr; 3 | use tokio::net::TcpListener; 4 | use tokio::task::JoinHandle; 5 | 6 | /// A PING PONG test without message provided. 7 | /// It should return "PONG". 8 | #[tokio::test] 9 | async fn ping_pong_without_message() { 10 | let (addr, _) = start_server().await; 11 | let mut client = client::connect(addr).await.unwrap(); 12 | 13 | let pong = client.ping(None).await.unwrap(); 14 | assert_eq!(b"PONG", &pong[..]); 15 | } 16 | 17 | /// A PING PONG test with message provided. 18 | /// It should return the message. 19 | #[tokio::test] 20 | async fn ping_pong_with_message() { 21 | let (addr, _) = start_server().await; 22 | let mut client = client::connect(addr).await.unwrap(); 23 | 24 | let pong = client.ping(Some("你好世界".to_string())).await.unwrap(); 25 | assert_eq!("你好世界".as_bytes(), &pong[..]); 26 | } 27 | 28 | /// A basic "hello world" style test. A server instance is started in a 29 | /// background task. A client instance is then established and set and get 30 | /// commands are sent to the server. The response is then evaluated 31 | #[tokio::test] 32 | async fn key_value_get_set() { 33 | let (addr, _) = start_server().await; 34 | 35 | let mut client = client::connect(addr).await.unwrap(); 36 | client.set("hello", "world".into()).await.unwrap(); 37 | 38 | let value = client.get("hello").await.unwrap().unwrap(); 39 | assert_eq!(b"world", &value[..]) 40 | } 41 | 42 | /// similar to the "hello world" style test, But this time 43 | /// a single channel subscription will be tested instead 44 | #[tokio::test] 45 | async fn receive_message_subscribed_channel() { 46 | let (addr, _) = start_server().await; 47 | 48 | let client = client::connect(addr).await.unwrap(); 49 | let mut subscriber = client.subscribe(vec!["hello".into()]).await.unwrap(); 50 | 51 | tokio::spawn(async move { 52 | let mut client = client::connect(addr).await.unwrap(); 53 | client.publish("hello", "world".into()).await.unwrap() 54 | }); 55 | 56 | let message = subscriber.next_message().await.unwrap().unwrap(); 57 | assert_eq!("hello", &message.channel); 58 | assert_eq!(b"world", &message.content[..]) 59 | } 60 | 61 | /// test that a client gets messages from multiple subscribed channels 62 | #[tokio::test] 63 | async fn receive_message_multiple_subscribed_channels() { 64 | let (addr, _) = start_server().await; 65 | 66 | let client = client::connect(addr).await.unwrap(); 67 | let mut subscriber = client 68 | .subscribe(vec!["hello".into(), "world".into()]) 69 | .await 70 | .unwrap(); 71 | 72 | tokio::spawn(async move { 73 | let mut client = client::connect(addr).await.unwrap(); 74 | client.publish("hello", "world".into()).await.unwrap() 75 | }); 76 | 77 | let message1 = subscriber.next_message().await.unwrap().unwrap(); 78 | assert_eq!("hello", &message1.channel); 79 | assert_eq!(b"world", &message1.content[..]); 80 | 81 | tokio::spawn(async move { 82 | let mut client = client::connect(addr).await.unwrap(); 83 | client.publish("world", "howdy?".into()).await.unwrap() 84 | }); 85 | 86 | let message2 = subscriber.next_message().await.unwrap().unwrap(); 87 | assert_eq!("world", &message2.channel); 88 | assert_eq!(b"howdy?", &message2.content[..]) 89 | } 90 | 91 | /// test that a client accurately removes its own subscribed chanel list 92 | /// when unsubscribing to all subscribed channels by submitting an empty vec 93 | #[tokio::test] 94 | async fn unsubscribes_from_channels() { 95 | let (addr, _) = start_server().await; 96 | 97 | let client = client::connect(addr).await.unwrap(); 98 | let mut subscriber = client 99 | .subscribe(vec!["hello".into(), "world".into()]) 100 | .await 101 | .unwrap(); 102 | 103 | subscriber.unsubscribe(&[]).await.unwrap(); 104 | assert_eq!(subscriber.get_subscribed().len(), 0); 105 | } 106 | 107 | async fn start_server() -> (SocketAddr, JoinHandle<()>) { 108 | let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); 109 | let addr = listener.local_addr().unwrap(); 110 | 111 | let handle = tokio::spawn(async move { server::run(listener, tokio::signal::ctrl_c()).await }); 112 | 113 | (addr, handle) 114 | } 115 | -------------------------------------------------------------------------------- /tests/server.rs: -------------------------------------------------------------------------------- 1 | use std::net::SocketAddr; 2 | 3 | use tokio::io::{AsyncReadExt, AsyncWriteExt}; 4 | use tokio::net::{TcpListener, TcpStream}; 5 | use tokio::time::{self, Duration}; 6 | 7 | use mini_redis::server; 8 | 9 | /// A basic "hello world" style test. A server instance is started in a 10 | /// background task. A client TCP connection is then established and raw redis 11 | /// commands are sent to the server. The response is evaluated at the byte 12 | /// level. 13 | #[tokio::test] 14 | async fn key_value_get_set() { 15 | let addr = start_server().await; 16 | 17 | // Establish a connection to the server 18 | let mut stream = TcpStream::connect(addr).await.unwrap(); 19 | 20 | // Get a key, data is missing 21 | stream 22 | .write_all(b"*2\r\n$3\r\nGET\r\n$5\r\nhello\r\n") 23 | .await 24 | .unwrap(); 25 | 26 | // Read nil response 27 | let mut response = [0; 5]; 28 | stream.read_exact(&mut response).await.unwrap(); 29 | assert_eq!(b"$-1\r\n", &response); 30 | 31 | // Set a key 32 | stream 33 | .write_all(b"*3\r\n$3\r\nSET\r\n$5\r\nhello\r\n$5\r\nworld\r\n") 34 | .await 35 | .unwrap(); 36 | 37 | // Read OK 38 | let mut response = [0; 5]; 39 | stream.read_exact(&mut response).await.unwrap(); 40 | assert_eq!(b"+OK\r\n", &response); 41 | 42 | // Get the key, data is present 43 | stream 44 | .write_all(b"*2\r\n$3\r\nGET\r\n$5\r\nhello\r\n") 45 | .await 46 | .unwrap(); 47 | 48 | // Shutdown the write half 49 | stream.shutdown().await.unwrap(); 50 | 51 | // Read "world" response 52 | let mut response = [0; 11]; 53 | stream.read_exact(&mut response).await.unwrap(); 54 | assert_eq!(b"$5\r\nworld\r\n", &response); 55 | 56 | // Receive `None` 57 | assert_eq!(0, stream.read(&mut response).await.unwrap()); 58 | } 59 | 60 | /// Similar to the basic key-value test, however, this time timeouts will be 61 | /// tested. This test demonstrates how to test time related behavior. 62 | /// 63 | /// When writing tests, it is useful to remove sources of non-determinism. Time 64 | /// is a source of non-determinism. Here, we "pause" time using the 65 | /// `time::pause()` function. This function is available with the `test-util` 66 | /// feature flag. This allows us to deterministically control how time appears 67 | /// to advance to the application. 68 | #[tokio::test] 69 | async fn key_value_timeout() { 70 | let addr = start_server().await; 71 | 72 | // Establish a connection to the server 73 | let mut stream = TcpStream::connect(addr).await.unwrap(); 74 | 75 | // Set a key 76 | stream 77 | .write_all( 78 | b"*5\r\n$3\r\nSET\r\n$5\r\nhello\r\n$5\r\nworld\r\n\ 79 | +EX\r\n:1\r\n", 80 | ) 81 | .await 82 | .unwrap(); 83 | 84 | let mut response = [0; 5]; 85 | 86 | // Read OK 87 | stream.read_exact(&mut response).await.unwrap(); 88 | 89 | assert_eq!(b"+OK\r\n", &response); 90 | 91 | // Get the key, data is present 92 | stream 93 | .write_all(b"*2\r\n$3\r\nGET\r\n$5\r\nhello\r\n") 94 | .await 95 | .unwrap(); 96 | 97 | // Read "world" response 98 | let mut response = [0; 11]; 99 | 100 | stream.read_exact(&mut response).await.unwrap(); 101 | 102 | assert_eq!(b"$5\r\nworld\r\n", &response); 103 | 104 | // Wait for the key to expire 105 | time::sleep(Duration::from_secs(1)).await; 106 | 107 | // Get a key, data is missing 108 | stream 109 | .write_all(b"*2\r\n$3\r\nGET\r\n$5\r\nhello\r\n") 110 | .await 111 | .unwrap(); 112 | 113 | // Read nil response 114 | let mut response = [0; 5]; 115 | 116 | stream.read_exact(&mut response).await.unwrap(); 117 | 118 | assert_eq!(b"$-1\r\n", &response); 119 | } 120 | 121 | #[tokio::test] 122 | async fn pub_sub() { 123 | let addr = start_server().await; 124 | 125 | let mut publisher = TcpStream::connect(addr).await.unwrap(); 126 | 127 | // Publish a message, there are no subscribers yet so the server will 128 | // return `0`. 129 | publisher 130 | .write_all(b"*3\r\n$7\r\nPUBLISH\r\n$5\r\nhello\r\n$5\r\nworld\r\n") 131 | .await 132 | .unwrap(); 133 | 134 | let mut response = [0; 4]; 135 | publisher.read_exact(&mut response).await.unwrap(); 136 | assert_eq!(b":0\r\n", &response); 137 | 138 | // Create a subscriber. This subscriber will only subscribe to the `hello` 139 | // channel. 140 | let mut sub1 = TcpStream::connect(addr).await.unwrap(); 141 | sub1.write_all(b"*2\r\n$9\r\nSUBSCRIBE\r\n$5\r\nhello\r\n") 142 | .await 143 | .unwrap(); 144 | 145 | // Read the subscribe response 146 | let mut response = [0; 34]; 147 | sub1.read_exact(&mut response).await.unwrap(); 148 | assert_eq!( 149 | &b"*3\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n:1\r\n"[..], 150 | &response[..] 151 | ); 152 | 153 | // Publish a message, there now is a subscriber 154 | publisher 155 | .write_all(b"*3\r\n$7\r\nPUBLISH\r\n$5\r\nhello\r\n$5\r\nworld\r\n") 156 | .await 157 | .unwrap(); 158 | 159 | let mut response = [0; 4]; 160 | publisher.read_exact(&mut response).await.unwrap(); 161 | assert_eq!(b":1\r\n", &response); 162 | 163 | // The first subscriber received the message 164 | let mut response = [0; 39]; 165 | sub1.read_exact(&mut response).await.unwrap(); 166 | assert_eq!( 167 | &b"*3\r\n$7\r\nmessage\r\n$5\r\nhello\r\n$5\r\nworld\r\n"[..], 168 | &response[..] 169 | ); 170 | 171 | // Create a second subscriber 172 | // 173 | // This subscriber will be subscribed to both `hello` and `foo` 174 | let mut sub2 = TcpStream::connect(addr).await.unwrap(); 175 | sub2.write_all(b"*3\r\n$9\r\nSUBSCRIBE\r\n$5\r\nhello\r\n$3\r\nfoo\r\n") 176 | .await 177 | .unwrap(); 178 | 179 | // Read the subscribe response 180 | let mut response = [0; 34]; 181 | sub2.read_exact(&mut response).await.unwrap(); 182 | assert_eq!( 183 | &b"*3\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n:1\r\n"[..], 184 | &response[..] 185 | ); 186 | let mut response = [0; 32]; 187 | sub2.read_exact(&mut response).await.unwrap(); 188 | assert_eq!( 189 | &b"*3\r\n$9\r\nsubscribe\r\n$3\r\nfoo\r\n:2\r\n"[..], 190 | &response[..] 191 | ); 192 | 193 | // Publish another message on `hello`, there are two subscribers 194 | publisher 195 | .write_all(b"*3\r\n$7\r\nPUBLISH\r\n$5\r\nhello\r\n$5\r\njazzy\r\n") 196 | .await 197 | .unwrap(); 198 | 199 | let mut response = [0; 4]; 200 | publisher.read_exact(&mut response).await.unwrap(); 201 | assert_eq!(b":2\r\n", &response); 202 | 203 | // Publish a message on `foo`, there is only one subscriber 204 | publisher 205 | .write_all(b"*3\r\n$7\r\nPUBLISH\r\n$3\r\nfoo\r\n$3\r\nbar\r\n") 206 | .await 207 | .unwrap(); 208 | 209 | let mut response = [0; 4]; 210 | publisher.read_exact(&mut response).await.unwrap(); 211 | assert_eq!(b":1\r\n", &response); 212 | 213 | // The first subscriber received the message 214 | let mut response = [0; 39]; 215 | sub1.read_exact(&mut response).await.unwrap(); 216 | assert_eq!( 217 | &b"*3\r\n$7\r\nmessage\r\n$5\r\nhello\r\n$5\r\njazzy\r\n"[..], 218 | &response[..] 219 | ); 220 | 221 | // The second subscriber received the message 222 | let mut response = [0; 39]; 223 | sub2.read_exact(&mut response).await.unwrap(); 224 | assert_eq!( 225 | &b"*3\r\n$7\r\nmessage\r\n$5\r\nhello\r\n$5\r\njazzy\r\n"[..], 226 | &response[..] 227 | ); 228 | 229 | // The first subscriber **did not** receive the second message 230 | let mut response = [0; 1]; 231 | time::timeout(Duration::from_millis(100), sub1.read(&mut response)) 232 | .await 233 | .unwrap_err(); 234 | 235 | // The second subscriber **did** receive the message 236 | let mut response = [0; 35]; 237 | sub2.read_exact(&mut response).await.unwrap(); 238 | assert_eq!( 239 | &b"*3\r\n$7\r\nmessage\r\n$3\r\nfoo\r\n$3\r\nbar\r\n"[..], 240 | &response[..] 241 | ); 242 | } 243 | 244 | #[tokio::test] 245 | async fn manage_subscription() { 246 | let addr = start_server().await; 247 | 248 | let mut publisher = TcpStream::connect(addr).await.unwrap(); 249 | 250 | // Create a subscriber 251 | let mut sub = TcpStream::connect(addr).await.unwrap(); 252 | sub.write_all(b"*2\r\n$9\r\nSUBSCRIBE\r\n$5\r\nhello\r\n") 253 | .await 254 | .unwrap(); 255 | 256 | // Read the subscribe response 257 | let mut response = [0; 34]; 258 | sub.read_exact(&mut response).await.unwrap(); 259 | assert_eq!( 260 | &b"*3\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n:1\r\n"[..], 261 | &response[..] 262 | ); 263 | 264 | // Update subscription to add `foo` 265 | sub.write_all(b"*2\r\n$9\r\nSUBSCRIBE\r\n$3\r\nfoo\r\n") 266 | .await 267 | .unwrap(); 268 | 269 | let mut response = [0; 32]; 270 | sub.read_exact(&mut response).await.unwrap(); 271 | assert_eq!( 272 | &b"*3\r\n$9\r\nsubscribe\r\n$3\r\nfoo\r\n:2\r\n"[..], 273 | &response[..] 274 | ); 275 | 276 | // Update subscription to remove `hello` 277 | sub.write_all(b"*2\r\n$11\r\nUNSUBSCRIBE\r\n$5\r\nhello\r\n") 278 | .await 279 | .unwrap(); 280 | 281 | let mut response = [0; 37]; 282 | sub.read_exact(&mut response).await.unwrap(); 283 | assert_eq!( 284 | &b"*3\r\n$11\r\nunsubscribe\r\n$5\r\nhello\r\n:1\r\n"[..], 285 | &response[..] 286 | ); 287 | 288 | // Publish a message to `hello` and then a message to `foo` 289 | publisher 290 | .write_all(b"*3\r\n$7\r\nPUBLISH\r\n$5\r\nhello\r\n$5\r\nworld\r\n") 291 | .await 292 | .unwrap(); 293 | let mut response = [0; 4]; 294 | publisher.read_exact(&mut response).await.unwrap(); 295 | assert_eq!(b":0\r\n", &response); 296 | 297 | publisher 298 | .write_all(b"*3\r\n$7\r\nPUBLISH\r\n$3\r\nfoo\r\n$3\r\nbar\r\n") 299 | .await 300 | .unwrap(); 301 | let mut response = [0; 4]; 302 | publisher.read_exact(&mut response).await.unwrap(); 303 | assert_eq!(b":1\r\n", &response); 304 | 305 | // Receive the message 306 | // The second subscriber **did** receive the message 307 | let mut response = [0; 35]; 308 | sub.read_exact(&mut response).await.unwrap(); 309 | assert_eq!( 310 | &b"*3\r\n$7\r\nmessage\r\n$3\r\nfoo\r\n$3\r\nbar\r\n"[..], 311 | &response[..] 312 | ); 313 | 314 | // No more messages 315 | let mut response = [0; 1]; 316 | time::timeout(Duration::from_millis(100), sub.read(&mut response)) 317 | .await 318 | .unwrap_err(); 319 | 320 | // Unsubscribe from all channels 321 | sub.write_all(b"*1\r\n$11\r\nunsubscribe\r\n") 322 | .await 323 | .unwrap(); 324 | 325 | let mut response = [0; 35]; 326 | sub.read_exact(&mut response).await.unwrap(); 327 | assert_eq!( 328 | &b"*3\r\n$11\r\nunsubscribe\r\n$3\r\nfoo\r\n:0\r\n"[..], 329 | &response[..] 330 | ); 331 | } 332 | 333 | // In this case we test that server Responds with an Error message if a client 334 | // sends an unknown command 335 | #[tokio::test] 336 | async fn send_error_unknown_command() { 337 | let addr = start_server().await; 338 | 339 | // Establish a connection to the server 340 | let mut stream = TcpStream::connect(addr).await.unwrap(); 341 | 342 | // Get a key, data is missing 343 | stream 344 | .write_all(b"*2\r\n$3\r\nFOO\r\n$5\r\nhello\r\n") 345 | .await 346 | .unwrap(); 347 | 348 | let mut response = [0; 28]; 349 | 350 | stream.read_exact(&mut response).await.unwrap(); 351 | 352 | assert_eq!(b"-err unknown command \'foo\'\r\n", &response); 353 | } 354 | 355 | // In this case we test that server Responds with an Error message if a client 356 | // sends an GET or SET command after a SUBSCRIBE 357 | #[tokio::test] 358 | async fn send_error_get_set_after_subscribe() { 359 | let addr = start_server().await; 360 | 361 | let mut stream = TcpStream::connect(addr).await.unwrap(); 362 | 363 | // send SUBSCRIBE command 364 | stream 365 | .write_all(b"*2\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n") 366 | .await 367 | .unwrap(); 368 | 369 | let mut response = [0; 34]; 370 | 371 | stream.read_exact(&mut response).await.unwrap(); 372 | 373 | assert_eq!( 374 | &b"*3\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n:1\r\n"[..], 375 | &response[..] 376 | ); 377 | 378 | stream 379 | .write_all(b"*3\r\n$3\r\nSET\r\n$5\r\nhello\r\n$5\r\nworld\r\n") 380 | .await 381 | .unwrap(); 382 | 383 | let mut response = [0; 28]; 384 | 385 | stream.read_exact(&mut response).await.unwrap(); 386 | assert_eq!(b"-err unknown command \'set\'\r\n", &response); 387 | 388 | stream 389 | .write_all(b"*2\r\n$3\r\nGET\r\n$5\r\nhello\r\n") 390 | .await 391 | .unwrap(); 392 | 393 | let mut response = [0; 28]; 394 | 395 | stream.read_exact(&mut response).await.unwrap(); 396 | assert_eq!(b"-err unknown command \'get\'\r\n", &response); 397 | } 398 | 399 | async fn start_server() -> SocketAddr { 400 | let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); 401 | let addr = listener.local_addr().unwrap(); 402 | 403 | tokio::spawn(async move { server::run(listener, tokio::signal::ctrl_c()).await }); 404 | 405 | addr 406 | } 407 | --------------------------------------------------------------------------------