From 8d304a2fe6352425777f96ba4edbced21049a4c6 Mon Sep 17 00:00:00 2001 From: bisho Date: Fri, 17 Nov 2023 15:40:19 +0100 Subject: [PATCH] Rust python module with helpers for meta-memcache (#1) * Rust python module with helpers for meta-memcache ## Motivation / Description Implementing this in rust is simple and benchmarking the rust cost gives 20-50 ns runtime. Unfortunately a lot of time is spent in the bridge between python and rust, so a call to rust code is not cheap, yet this implementation speeds up some meta-memcache code significantly. * split tests, u8 for response type * More safety parsing the header and simplify code * rename module --- .github/workflows/CI.yml | 109 +++++++++++ .gitignore | 72 ++++++++ Cargo.lock | 306 +++++++++++++++++++++++++++++++ Cargo.toml | 17 ++ meta_memcache_socket.pyi | 174 ++++++++++++++++++ pyproject.toml | 16 ++ src/constants.rs | 20 ++ src/impl_build_cmd.rs | 71 ++++++++ src/impl_build_cmd_tests.rs | 194 ++++++++++++++++++++ src/impl_parse_header.rs | 51 ++++++ src/impl_parse_header_tests.rs | 246 +++++++++++++++++++++++++ src/lib.rs | 172 ++++++++++++++++++ src/request_flags.rs | 322 +++++++++++++++++++++++++++++++++ src/response_flags.rs | 248 +++++++++++++++++++++++++ 14 files changed, 2018 insertions(+) create mode 100644 .github/workflows/CI.yml create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 meta_memcache_socket.pyi create mode 100644 pyproject.toml create mode 100644 src/constants.rs create mode 100644 src/impl_build_cmd.rs create mode 100644 src/impl_build_cmd_tests.rs create mode 100644 src/impl_parse_header.rs create mode 100644 src/impl_parse_header_tests.rs create mode 100644 src/lib.rs create mode 100644 src/request_flags.rs create mode 100644 src/response_flags.rs diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml new file mode 100644 index 0000000..e0fdb3a --- /dev/null +++ b/.github/workflows/CI.yml @@ -0,0 +1,109 @@ +# This file is autogenerated by maturin v1.3.0 +# To update, run +# +# maturin generate-ci github +# +name: CI + +on: + push: + branches: + - main + - master + tags: + - '*' + pull_request: + workflow_dispatch: + +permissions: + contents: read + +jobs: + tests: + runs-on: ubuntu-latest + name: "Tests" + steps: + - uses: actions/checkout@v3 + - uses: actions-rs/toolchain@v1 + with: + toolchain: stable + - uses: actions-rs/cargo@v1 + with: + command: test + + linux: + runs-on: ubuntu-latest + strategy: + matrix: + target: [x86_64, x86, aarch64, armv7] + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v4 + with: + python-version: '3.11' + - name: Build wheels + uses: PyO3/maturin-action@v1 + with: + target: ${{ matrix.target }} + args: --release --out dist --find-interpreter + sccache: 'true' + manylinux: auto + - name: Upload wheels + uses: actions/upload-artifact@v3 + with: + name: wheels + path: dist + + macos: + runs-on: macos-latest + strategy: + matrix: + target: [x86_64, aarch64] + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v4 + with: + python-version: '3.11' + - name: Build wheels + uses: PyO3/maturin-action@v1 + with: + target: ${{ matrix.target }} + args: --release --out dist --find-interpreter + sccache: 'true' + - name: Upload wheels + uses: actions/upload-artifact@v3 + with: + name: wheels + path: dist + + sdist: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Build sdist + uses: PyO3/maturin-action@v1 + with: + command: sdist + args: --out dist + - name: Upload sdist + uses: actions/upload-artifact@v3 + with: + name: wheels + path: dist + + release: + name: Release + runs-on: ubuntu-latest + if: "startsWith(github.ref, 'refs/tags/')" + needs: [linux, macos, sdist] + steps: + - uses: actions/download-artifact@v3 + with: + name: wheels + - name: Publish to PyPI + uses: PyO3/maturin-action@v1 + env: + MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }} + with: + command: upload + args: --non-interactive --skip-existing * diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..af3ca5e --- /dev/null +++ b/.gitignore @@ -0,0 +1,72 @@ +/target + +# Byte-compiled / optimized / DLL files +__pycache__/ +.pytest_cache/ +*.py[cod] + +# C extensions +*.so + +# Distribution / packaging +.Python +.venv/ +env/ +bin/ +build/ +develop-eggs/ +dist/ +eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +include/ +man/ +venv/ +*.egg-info/ +.installed.cfg +*.egg + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt +pip-selfcheck.json + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.cache +nosetests.xml +coverage.xml + +# Translations +*.mo + +# Mr Developer +.mr.developer.cfg +.project +.pydevproject + +# Rope +.ropeproject + +# Django stuff: +*.log +*.pot + +.DS_Store + +# Sphinx documentation +docs/_build/ + +# PyCharm +.idea/ + +# VSCode +.vscode/ + +# Pyenv +.python-version \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..7d23698 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,306 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + +[[package]] +name = "autocfg" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" + +[[package]] +name = "base64" +version = "0.21.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35636a1494ede3b646cc98f74f8e62c773a38a659ebc777a2cf26b9b74171df9" + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "heck" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" + +[[package]] +name = "indoc" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e186cfbae8084e513daff4240b4797e342f988cecda4fb6c939150f96315fd8" + +[[package]] +name = "libc" +version = "0.2.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89d92a4743f9a61002fae18374ed11e7973f530cb3a3255fb354818118b2203c" + +[[package]] +name = "lock_api" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45" +dependencies = [ + "autocfg", + "scopeguard", +] + +[[package]] +name = "memoffset" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c" +dependencies = [ + "autocfg", +] + +[[package]] +name = "meta-memcache-socket" +version = "0.1.0" +dependencies = [ + "atoi", + "base64", + "pyo3", +] + +[[package]] +name = "num-traits" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d" + +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-targets", +] + +[[package]] +name = "proc-macro2" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "134c189feb4956b20f6f547d2cf727d4c0fe06722b20a0eec87ed445a97f92da" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04e8453b658fe480c3e70c8ed4e3d3ec33eb74988bd186561b0cc66b85c3bc4b" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "parking_lot", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a96fe70b176a89cff78f2fa7b3c930081e163d5379b4dcdf993e3ae29ca662e5" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "214929900fd25e6604661ed9cf349727c8920d47deff196c4e28165a6ef2a96b" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dac53072f717aa1bfa4db832b39de8c875b7c7af4f4a6fe93cdbf9264cf8383b" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7774b5a8282bd4f25f803b1f0d945120be959a36c72e08e7cd031c792fdfd424" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5267fca4496028628a95160fc423a33e8b2e6af8a5302579e322e4b520293cae" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "redox_syscall" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +dependencies = [ + "bitflags", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "smallvec" +version = "1.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "942b4a808e05215192e39f4ab80813e599068285906cc91aa64f923db842bd5a" + +[[package]] +name = "syn" +version = "2.0.39" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23e78b90f2fcf45d3e842032ce32e3f2d1545ba6636271dcbf24fa306d87be7a" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.12.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14c39fd04924ca3a864207c66fc2cd7d22d7c016007f9ce846cbb9326331930a" + +[[package]] +name = "unicode-ident" +version = "1.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" + +[[package]] +name = "unindent" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" + +[[package]] +name = "windows-targets" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" + +[[package]] +name = "windows_i686_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" + +[[package]] +name = "windows_i686_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..6f4a4ad --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "meta-memcache-socket" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[lib] +name = "meta_memcache_socket" +crate-type = ["cdylib"] + +[dependencies] +atoi = "2.0.0" +base64 = "0.21.5" + +[dependencies.pyo3] +version = "0.20.0" +features = ["extension-module"] diff --git a/meta_memcache_socket.pyi b/meta_memcache_socket.pyi new file mode 100644 index 0000000..532c4ab --- /dev/null +++ b/meta_memcache_socket.pyi @@ -0,0 +1,174 @@ +from typing import Optional, Tuple, Union + +RESPONSE_VALUE = 1 # VALUE (VA) +RESPONSE_SUCCESS = 2 # SUCCESS (OK or HD) +RESPONSE_NOT_STORED = 3 # NOT_STORED (NS) +RESPONSE_CONFLICT = 4 # CONFLICT (EX) +RESPONSE_MISS = 5 # MISS (EN or NF) +RESPONSE_NOOP = 100 # NOOP (MN) + +# Set modes +# E "add" command. LRU bump and return NS if item exists. Else add. +SET_MODE_ADD = 69 +# A "append" command. If item exists, append the new value to its data. +SET_MODE_APPEND = 65 # 'A' +# P "prepend" command. If item exists, prepend the new value to its data. +SET_MODE_PREPEND = 80 # 'P' +# R "replace" command. Set only if item already exists. +SET_MODE_REPLACE = 82 # 'R' +# S "set" command. The default mode, added for completeness. +SET_MODE_SET = 83 # 'S' +# Arithmetic modes +# + "increment" +MA_MODE_INC = 43 +# - "decrement" +MA_MODE_DEC = 45 + +class RequestFlags: + """ + A class representing the flags for a meta-protocol request + + * no_reply: Set to True if the server should not send a response + * return_client_flag: Set to True if the server should return the client flag + * return_cas_token: Set to True if the server should return the CAS token + * return_value: Set to True if the server should return the value (Default) + * return_ttl: Set to True if the server should return the TTL + * return_size: Set to True if the server should return the size (useful + if when paired with return_value=False, to get the size of the value) + * return_last_access: Set to True if the server should return the last access time + * return_fetched: Set to True if the server should return the fetched flag + * return_key: Set to True if the server should return the key in the response + * no_update_lru: Set to True if the server should not update the LRU on this access + * mark_stale: Set to True if the server should mark the value as stale + * cache_ttl: The TTL to set on the key + * recache_ttl: The TTL to use for recache policy + * vivify_on_miss_ttl: The TTL to use when vivifying a value on a miss + * client_flag: The client flag to store along the value (Useful to store value type, compression, etc) + * ma_initial_value: For arithmetic operations, the initial value to use (if the key does not exist) + * ma_delta_value: For arithmetic operations, the delta value to use + * cas_token: The CAS token to use when storing the value in the cache + * opaque: The opaque flag (will be echoed back in the response) + * mode: The mode to use when storing the value in the cache. See SET_MODE_* and MA_MODE_* constants + """ + + no_reply: bool + return_client_flag: bool + return_cas_token: bool + return_value: bool + return_ttl: bool + return_size: bool + return_last_access: bool + return_fetched: bool + return_key: bool + no_update_lru: bool + mark_stale: bool + cache_ttl: Optional[int] + recache_ttl: Optional[int] + vivify_on_miss_ttl: Optional[int] + client_flag: Optional[int] + ma_initial_value: Optional[int] + ma_delta_value: Optional[int] + cas_token: Optional[int] + opaque: Optional[bytes] + mode: Optional[int] + + def __init__( + *, + no_reply: bool = False, + return_client_flag: bool = True, + return_cas_token: bool = False, + return_value=True, + return_ttl: bool = False, + return_size: bool = False, + return_last_access: bool = False, + return_fetched: bool = False, + return_key: bool = False, + no_update_lru: bool = False, + mark_stale: bool = False, + cache_ttl: Optional[int] = None, + recache_ttl: Optional[int] = None, + vivify_on_miss_ttl: Optional[int] = None, + client_flag: Optional[int] = None, + ma_initial_value: Optional[int] = None, + ma_delta_value: Optional[int] = None, + cas_token: Optional[int] = None, + opaque: Optional[bytes] = None, + mode: Optional[int] = None, + ) -> None: ... + def copy(self) -> "RequestFlags": ... + def to_bytes(self) -> bytes: ... + +class ResponseFlags: + """ + A class representing the flags for a meta-protocol response + + * cas_token: Compare-And-Swap token (integer value) or None if not returned + * fetched: + - True if fetched since being set + - False if not fetched since being set + - None if the server di not return this flag info + * last_access: time in seconds since last access (integer value) or None if not returned + * ttl: time in seconds until the value expires (integer value) or None if not returned + - The special value -1 represents if the key will never expire + * client_flag: integer value or None if not returned + * win: + - True if the client won the right to repopulate + - False if the client lost the right to repopulate + - None if the server did not return a win/lose flag + * stale: True if the value is stale, False otherwise + * real_size: integer value or None if not returned + * opaque flag: bytes value or None if not returned + """ + + cas_token: Optional[int] + fetched: Optional[bool] + last_access: Optional[int] + ttl: Optional[int] + client_flag: Optional[int] + win: Optional[bool] + stale: bool + real_size: Optional[int] + opaque: Optional[bytes] + + def __init__( + *, + cas_token=None, + fetched=None, + last_access=None, + ttl=None, + client_flag=None, + win=None, + stale=False, + size=None, + opaque=None, + ) -> None: ... + +def parse_header( + buffer: Union[memoryview, bytes, bytearray], + start: int, + end: int, +) -> Optional[Tuple[int, Optional[int], Optional[int], Optional[ResponseFlags]]]: + """ + Parse a memcache meta-protocol header from a buffer + + :param buffer: The buffer to parse + :param start: The starting point in the buffer + :param end: The end of the data read into the buffer + """ + ... + +def build_cmd( + cmd: bytes, + key: bytes, + size: Optional[int] = None, + flags: Optional[RequestFlags] = None, + legacy_size_format: bool = False, +) -> bytes: + """ + Build a memcache meta-protocol command + + :param cmd: The command to send + :param key: The key to use + :param flags: The flags to use + """ + ... diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..6f4bd7c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,16 @@ +[build-system] +requires = ["maturin>=1.3,<2.0"] +build-backend = "maturin" + +[project] +name = "meta-memcache-socket" +requires-python = ">=3.7" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Python :: Implementation :: PyPy", +] +dynamic = ["version"] + +[tool.maturin] +features = ["pyo3/extension-module"] diff --git a/src/constants.rs b/src/constants.rs new file mode 100644 index 0000000..bc0b12a --- /dev/null +++ b/src/constants.rs @@ -0,0 +1,20 @@ +pub const RESPONSE_VALUE: u8 = 1; // VALUE (VA) +pub const RESPONSE_SUCCESS: u8 = 2; // SUCCESS (OK or HD) +pub const RESPONSE_NOT_STORED: u8 = 3; // NOT_STORED (NS) +pub const RESPONSE_CONFLICT: u8 = 4; // CONFLICT (EX) +pub const RESPONSE_MISS: u8 = 5; // MISS (EN or NF) +pub const RESPONSE_NOOP: u8 = 100; // NOOP (MN) + +// Set modes: +// E: "add" command. LRU bump and return NS if item exists. Else add. +// A: "append" command. If item exists, append the new value to its data. +// P: "prepend" command. If item exists, prepend the new value to its data. +// R: "replace" command. Set only if item already exists. +// S: "set" command. The default mode, added for completeness. +pub const SET_MODE_ADD: u8 = 69; // 'E' +pub const SET_MODE_APPEND: u8 = 65; // 'A' +pub const SET_MODE_PREPEND: u8 = 80; // 'P' +pub const SET_MODE_REPLACE: u8 = 82; // 'R' +pub const SET_MODE_SET: u8 = 83; // 'S' +pub const MA_MODE_INC: u8 = 43; // '+' +pub const MA_MODE_DEC: u8 = 45; // '-' diff --git a/src/impl_build_cmd.rs b/src/impl_build_cmd.rs new file mode 100644 index 0000000..0409055 --- /dev/null +++ b/src/impl_build_cmd.rs @@ -0,0 +1,71 @@ +use base64::{engine::general_purpose, Engine as _}; + +use crate::RequestFlags; + +const MAX_KEY_LEN: usize = 250; +const MAX_BIN_KEY_LEN: usize = 187; // 250 * 3 / 4 due to b64 encoding + +pub fn impl_build_cmd( + cmd: &[u8], + key: &[u8], + size: Option, + request_flags: Option<&RequestFlags>, + legacy_size_format: bool, +) -> Option> { + if key.len() >= MAX_KEY_LEN { + // Key is too long + return None; + } + let mut binary = false; + for c in key.iter() { + if *c <= b' ' || *c > b'~' { + // Not ascii or containing spaces + binary = true; + break; + } + } + if binary && key.len() >= MAX_BIN_KEY_LEN { + // Key is too long + return None; + } + + // Build the command + let mut buf: Vec = Vec::new(); + + // Add CMD + buf.extend_from_slice(cmd); + buf.push(b' '); + + // Add key + if binary { + // If the key contains binary or spaces, it will be send in b64 + let result = general_purpose::STANDARD.encode(key); + buf.extend_from_slice(&result.as_bytes()); + } else { + // Otherwise, it will be send as is + buf.extend_from_slice(key); + } + + // Add size + if let Some(size) = size { + buf.push(b' '); + if legacy_size_format { + buf.push(b'S'); + } + buf.extend_from_slice(&size.to_string().as_bytes()); + } + + // Add request flags + if binary { + // If the key is binary, it will be send in b64. Adding the b flag will + // tell the server to decode it and store it as binary, saving memory. + buf.push(b' '); + buf.push(b'b'); + } + if let Some(request_flags) = request_flags { + request_flags.push_bytes(&mut buf); + } + buf.push(b'\r'); + buf.push(b'\n'); + Some(buf) +} diff --git a/src/impl_build_cmd_tests.rs b/src/impl_build_cmd_tests.rs new file mode 100644 index 0000000..86be468 --- /dev/null +++ b/src/impl_build_cmd_tests.rs @@ -0,0 +1,194 @@ +#[cfg(test)] +mod tests { + use crate::impl_build_cmd; + use crate::request_flags::RequestFlags; + + #[test] + fn test_impl_build_cmd_with_flags() { + let cmd = b"mg"; + let key = b"key"; + let request_flags = RequestFlags::new( + true, // no_reply + true, // return_client_flag + true, // return_cas_token + true, // return_value + true, // return_ttl + true, // return_size + true, // return_last_access + true, // return_fetched + true, // return_key + true, // no_update_lru + true, // mark_stale + Some(111), // cache_ttl + Some(222), // recache_ttl + Some(333), // vivify_on_miss_ttl + Some(444), // client_flag + Some(555), // ma_initial_value + Some(666), // ma_delta_value, + Some(777), // cas_token + Some(b"opaque".to_vec()), // opaque + Some(65 as u8), // mode + ); + + let result = impl_build_cmd(cmd, key, None, Some(&request_flags), false).unwrap(); + let string = String::from_utf8_lossy(&result); + println!("{:?}", string); + assert_eq!( + result, + b"mg key q f c v t s l h k u I T111 R222 N333 F444 J555 D666 C777 Oopaque MA\r\n" + ); + } + + #[test] + fn test_impl_build_cmd_no_flags() { + let cmd = b"mg"; + let key = b"key"; + let request_flags = RequestFlags::new( + false, // no_reply + false, // return_client_flag + false, // return_cas_token + false, // return_value + false, // return_ttl + false, // return_size + false, // return_last_access + false, // return_fetched + false, // return_key + false, // no_update_lru + false, // mark_stale + None, // cache_ttl + None, // recache_ttl + None, // vivify_on_miss_ttl + None, // client_flag + None, // ma_initial_value + None, // ma_delta_value, + None, // cas_token + None, // opaque + None, // mode + ); + + let result = impl_build_cmd(cmd, key, None, Some(&request_flags), false).unwrap(); + let string = String::from_utf8_lossy(&result); + println!("{:?}", string); + assert_eq!(result, b"mg key\r\n"); + } + + #[test] + fn test_impl_build_cmd_binary_key() { + let cmd = b"mg"; + let key = b"Key_with_binary\x00"; + let request_flags = RequestFlags::new( + false, // no_reply + false, // return_client_flag + false, // return_cas_token + false, // return_value + false, // return_ttl + false, // return_size + false, // return_last_access + false, // return_fetched + false, // return_key + false, // no_update_lru + false, // mark_stale + None, // cache_ttl + None, // recache_ttl + None, // vivify_on_miss_ttl + None, // client_flag + None, // ma_initial_value + None, // ma_delta_value, + None, // cas_token + None, // opaque + None, // mode + ); + + let result = impl_build_cmd(cmd, key, None, Some(&request_flags), false).unwrap(); + let string = String::from_utf8_lossy(&result); + println!("{:?}", string); + assert_eq!(result, b"mg S2V5X3dpdGhfYmluYXJ5AA== b\r\n"); + } + + #[test] + fn test_impl_build_cmd_key_with_spaces() { + let cmd = b"mg"; + let key = b"Key with spaces"; + let request_flags = RequestFlags::new( + false, // no_reply + false, // return_client_flag + false, // return_cas_token + false, // return_value + false, // return_ttl + false, // return_size + false, // return_last_access + false, // return_fetched + false, // return_key + false, // no_update_lru + false, // mark_stale + None, // cache_ttl + None, // recache_ttl + None, // vivify_on_miss_ttl + None, // client_flag + None, // ma_initial_value + None, // ma_delta_value, + None, // cas_token + None, // opaque + None, // mode + ); + + let result = impl_build_cmd(cmd, key, None, Some(&request_flags), false).unwrap(); + let string = String::from_utf8_lossy(&result); + println!("{:?}", string); + assert_eq!(result, b"mg S2V5IHdpdGggc3BhY2Vz b\r\n"); + } + + #[test] + fn test_impl_build_cmd_large_key() { + let cmd = b"mg"; + let key = &vec![b'X'; 250]; + let no_result = impl_build_cmd(cmd, key, None, None, false); + assert!(no_result.is_none()); + } + + #[test] + fn test_cmd_with_size() { + let cmd = b"ms"; + let key = b"key"; + let size = 123; + let request_flags = RequestFlags::new( + false, // no_reply + false, // return_client_flag + false, // return_cas_token + false, // return_value + false, // return_ttl + false, // return_size + false, // return_last_access + false, // return_fetched + false, // return_key + false, // no_update_lru + false, // mark_stale + Some(111), // cache_ttl + None, // recache_ttl + None, // vivify_on_miss_ttl + None, // client_flag + None, // ma_initial_value + None, // ma_delta_value, + None, // cas_token + None, // opaque + None, // mode + ); + + let result = impl_build_cmd(cmd, key, Some(size), Some(&request_flags), false).unwrap(); + let string = String::from_utf8_lossy(&result); + println!("{:?}", string); + assert_eq!(result, b"ms key 123 T111\r\n"); + } + + #[test] + fn test_cmd_with_legacy_size() { + let cmd = b"ms"; + let key = b"key"; + let size = 123; + + let result = impl_build_cmd(cmd, key, Some(size), None, true).unwrap(); + let string = String::from_utf8_lossy(&result); + println!("{:?}", string); + assert_eq!(result, b"ms key S123\r\n"); + } +} diff --git a/src/impl_parse_header.rs b/src/impl_parse_header.rs new file mode 100644 index 0000000..531d7b2 --- /dev/null +++ b/src/impl_parse_header.rs @@ -0,0 +1,51 @@ +use crate::{constants::*, ResponseFlags}; + +pub fn impl_parse_header( + data: &[u8], + start: usize, + end: usize, +) -> Option<(usize, Option, Option, Option)> { + if end - start < 4 { + return None; + } + let end = end.min(data.len()); + let mut n = start + 2; + while n < end - 1 { + if data[n] == b'\r' && data[n + 1] == b'\n' { + let endl_pos = n + 2; + match &data[start..start + 2] { + b"VA" => { + match ResponseFlags::from_value_header(&data[start..n]) { + Some((size, flags)) => { + return Some((endl_pos, Some(RESPONSE_VALUE), Some(size), Some(flags))); + } + None => { + return Some((endl_pos, None, None, None)); + } + }; + } + b"HD" | b"OK" => { + let flags = ResponseFlags::from_success_header(&data[start..n]); + return Some((endl_pos, Some(RESPONSE_SUCCESS), None, Some(flags))); + } + b"NS" => { + return Some((endl_pos, Some(RESPONSE_NOT_STORED), None, None)); + } + b"EX" => { + return Some((endl_pos, Some(RESPONSE_CONFLICT), None, None)); + } + b"EN" | b"NF" => { + return Some((endl_pos, Some(RESPONSE_MISS), None, None)); + } + b"MN" => { + return Some((endl_pos, Some(RESPONSE_NOOP), None, None)); + } + _ => { + return Some((endl_pos, None, None, None)); + } + } + } + n += 1; + } + None +} diff --git a/src/impl_parse_header_tests.rs b/src/impl_parse_header_tests.rs new file mode 100644 index 0000000..6e5c778 --- /dev/null +++ b/src/impl_parse_header_tests.rs @@ -0,0 +1,246 @@ +#[cfg(test)] +mod tests { + use crate::constants::*; + use crate::impl_parse_header; + + #[test] + fn test_no_crnl_in_buffer() { + let data = b"X\rX\nX"; + let no_result = impl_parse_header(data, 0, data.len()); + assert!(no_result.is_none()); + } + #[test] + fn test_value_response() { + let data = b"VA 1234 c1234567 h0 l1111 t2222 f1 Z s3333 MORE_SPACES_ARE_OK_TOO Ofooonly UNKNOWN FLAGS Ofoobar\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_VALUE)); + assert_eq!(size, Some(1234)); + assert!(flags.is_some()); + let flags = flags.unwrap(); + assert_eq!(flags.cas_token, Some(1234567)); + assert_eq!(flags.fetched, Some(false)); + assert_eq!(flags.last_access, Some(1111)); + assert_eq!(flags.ttl, Some(2222)); + assert_eq!(flags.client_flag, Some(1)); + assert_eq!(flags.win, Some(false)); + assert_eq!(flags.stale, false); + assert_eq!(flags.size, Some(3333)); + assert_eq!(flags.opaque, Some(b"foobar".to_vec())); + } + + #[test] + fn test_size_int_overflow() { + let data = b"VA 12345678901234567890 c123456789001234567890 l111 t12345678901234567890\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(end_pos, data.len()); + assert!(response_type.is_none()); + assert!(size.is_none()); + assert!(flags.is_none()); + } + + #[test] + fn test_flags_int_overflow() { + let data = b"VA 1234 c123456789001234567890 l111 t12345678901234567890\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_VALUE)); + assert_eq!(size, Some(1234)); + assert!(flags.is_some()); + let flags = flags.unwrap(); + assert!(flags.cas_token.is_none()); + assert!(flags.fetched.is_none()); + assert_eq!(flags.last_access, Some(111)); + assert!(flags.ttl.is_none()); + assert!(flags.client_flag.is_none()); + assert!(flags.win.is_none()); + assert_eq!(flags.stale, false); + assert!(flags.size.is_none()); + assert!(flags.opaque.is_none()); + } + + #[test] + fn test_bad_ttls() { + let data = b"VA 1234 c111 t\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_VALUE)); + assert_eq!(size, Some(1234)); + assert!(flags.is_some()); + let flags = flags.unwrap(); + assert_eq!(flags.cas_token, Some(111)); + assert!(flags.fetched.is_none()); + assert!(flags.last_access.is_none()); + assert!(flags.ttl.is_none()); + assert!(flags.client_flag.is_none()); + assert!(flags.win.is_none()); + assert_eq!(flags.stale, false); + assert!(flags.size.is_none()); + assert!(flags.opaque.is_none()); + let data = b"VA 1234 t-999 c111\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_VALUE)); + assert_eq!(size, Some(1234)); + assert!(flags.is_some()); + let flags = flags.unwrap(); + assert_eq!(flags.cas_token, Some(111)); + assert!(flags.fetched.is_none()); + assert!(flags.last_access.is_none()); + assert_eq!(flags.ttl, Some(-1)); + assert!(flags.client_flag.is_none()); + assert!(flags.win.is_none()); + assert_eq!(flags.stale, false); + assert!(flags.size.is_none()); + assert!(flags.opaque.is_none()); + let data = b"VA 1234 t- c111\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_VALUE)); + assert_eq!(size, Some(1234)); + assert!(flags.is_some()); + let flags = flags.unwrap(); + assert_eq!(flags.cas_token, Some(111)); + assert!(flags.fetched.is_none()); + assert!(flags.last_access.is_none()); + assert_eq!(flags.ttl, Some(-1)); + assert!(flags.client_flag.is_none()); + assert!(flags.win.is_none()); + assert_eq!(flags.stale, false); + assert!(flags.size.is_none()); + assert!(flags.opaque.is_none()); + } + + #[test] + fn test_value_response_no_flags() { + let data = b"VA 1234\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_VALUE)); + assert_eq!(size, Some(1234)); + assert!(flags.is_some()); + let flags = flags.unwrap(); + assert!(flags.cas_token.is_none()); + assert!(flags.fetched.is_none()); + assert!(flags.last_access.is_none()); + assert!(flags.ttl.is_none()); + assert!(flags.client_flag.is_none()); + assert!(flags.win.is_none()); + assert_eq!(flags.stale, false); + assert!(flags.size.is_none()); + assert!(flags.opaque.is_none()); + } + + #[test] + fn test_value_response_no_size() { + let data = b"VA c123\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert!(response_type.is_none()); + assert!(size.is_none()); + assert!(flags.is_none()); + } + + #[test] + fn test_success_reponse() { + let data = b"HD c1234567 h0 l1111 t-1 f1 X W s2222 Ofoobar UNKNOWN FLAGS\r\nOK\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len() - 4); + assert_eq!(response_type, Some(RESPONSE_SUCCESS)); + assert!(size.is_none()); + assert!(flags.is_some()); + let flags = flags.unwrap(); + assert_eq!(flags.cas_token, Some(1234567)); + assert_eq!(flags.fetched, Some(false)); + assert_eq!(flags.last_access, Some(1111)); + assert_eq!(flags.ttl, Some(-1)); + assert_eq!(flags.client_flag, Some(1)); + assert_eq!(flags.win, Some(true)); + assert_eq!(flags.stale, true); + assert_eq!(flags.size, Some(2222)); + assert_eq!(flags.opaque, Some(b"foobar".to_vec())); + let (end_pos, response_type, size, flags) = + impl_parse_header(data, data.len() - 4, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_SUCCESS)); + assert!(size.is_none()); + assert!(flags.is_some()); + let flags = flags.unwrap(); + assert!(flags.cas_token.is_none()); + assert!(flags.fetched.is_none()); + assert!(flags.last_access.is_none()); + assert!(flags.ttl.is_none()); + assert!(flags.client_flag.is_none()); + assert!(flags.win.is_none()); + assert_eq!(flags.stale, false); + assert!(flags.size.is_none()); + assert!(flags.opaque.is_none()); + } + + #[test] + fn test_not_stored_response() { + let data = b"NS\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_NOT_STORED)); + assert!(size.is_none()); + assert!(flags.is_none()); + } + #[test] + fn test_conflict_response() { + let data = b"EX\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_CONFLICT)); + assert!(size.is_none()); + assert!(flags.is_none()); + } + #[test] + fn test_miss_response() { + let data = b"EN\r\nNF\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, 4); // Only reads the first response header + assert_eq!(response_type, Some(RESPONSE_MISS)); + assert!(size.is_none()); + assert!(flags.is_none()); + let (end_pos, response_type, size, flags) = impl_parse_header(data, 4, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_MISS)); + assert!(size.is_none()); + assert!(flags.is_none()); + } + #[test] + fn test_noop_response() { + let data = b"MN\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert_eq!(response_type, Some(RESPONSE_NOOP)); + assert!(size.is_none()); + assert!(flags.is_none()); + } + + #[test] + fn test_unknown_response() { + let data = b"XX 33 c1 Z f1\r\n"; + let (end_pos, response_type, size, flags) = impl_parse_header(data, 0, data.len()).unwrap(); + assert_eq!(end_pos, data.len()); + assert!(response_type.is_none()); + assert!(size.is_none()); + assert!(flags.is_none()); + } + + #[test] + fn test_response_too_small() { + let data = b"X\r\n"; + let no_result = impl_parse_header(data, 0, data.len()); + assert!(no_result.is_none()); + } + + #[test] + fn test_end_is_out_of_bounds() { + let data = b"NOENDLINE"; + let no_result = impl_parse_header(data, 0, data.len() + 100); + assert!(no_result.is_none()); + } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..201e4e7 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,172 @@ +mod constants; +mod impl_build_cmd; +mod impl_build_cmd_tests; +mod impl_parse_header; +mod impl_parse_header_tests; +mod request_flags; +mod response_flags; +pub use constants::*; +use impl_build_cmd::impl_build_cmd; +use impl_parse_header::impl_parse_header; +pub use request_flags::RequestFlags; +pub use response_flags::ResponseFlags; + +use std::slice; + +use pyo3::buffer::PyBuffer; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +#[pyfunction] +#[pyo3( + signature = ( + buffer, + start, + end, + ), + text_signature = "(buffer: Union[memoryview, bytearray], start: int, end: int)", +)] +pub fn parse_header( + buffer: PyBuffer, + start: usize, + end: usize, +) -> PyResult, Option, Option)>> { + if end > buffer.len_bytes() { + return Err(pyo3::exceptions::PyValueError::new_err( + "End must be less than buffer length", + )); + } + let data = unsafe { slice::from_raw_parts(buffer.buf_ptr() as *const u8, buffer.len_bytes()) }; + Ok(impl_parse_header(data, start, end)) +} + +#[pyfunction] +#[pyo3( + signature = ( + cmd, + key, + size=None, + request_flags=None, + legacy_size_format=false, + ), + text_signature = "(cmd: bytes, key: bytes, size: Optional[int], request_flags: Optional[RequestFlags], legacy_size_format: bool = False)", +)] +pub fn build_cmd( + py: Python, + cmd: &[u8], + key: &[u8], + size: Option, + request_flags: Option<&RequestFlags>, + legacy_size_format: bool, +) -> PyResult> { + match impl_build_cmd(cmd, key, size, request_flags, legacy_size_format) { + Some(buf) => Ok(PyBytes::new(py, &buf).into()), + None => Err(pyo3::exceptions::PyValueError::new_err("Key is too long")), + } +} + +#[pyfunction] +#[pyo3( + signature = ( + key, + request_flags=None, + ), + text_signature = "(key: bytes, request_flags: Optional[RequestFlags])", +)] +pub fn build_meta_get( + py: Python, + key: &[u8], + request_flags: Option<&RequestFlags>, +) -> PyResult> { + match impl_build_cmd(b"mg", key, None, request_flags, false) { + Some(buf) => Ok(PyBytes::new(py, &buf).into()), + None => Err(pyo3::exceptions::PyValueError::new_err("Key is too long")), + } +} + +#[pyfunction] +#[pyo3( + signature = ( + key, + size, + request_flags=None, + legacy_size_format=false, + ), + text_signature = "(key: bytes, size: int, request_flags: Optional[RequestFlags], legacy_size_format: bool = False)", +)] +pub fn build_meta_set( + py: Python, + key: &[u8], + size: u32, + request_flags: Option<&RequestFlags>, + legacy_size_format: bool, +) -> PyResult> { + match impl_build_cmd(b"ms", key, Some(size), request_flags, legacy_size_format) { + Some(buf) => Ok(PyBytes::new(py, &buf).into()), + None => Err(pyo3::exceptions::PyValueError::new_err("Key is too long")), + } +} + +#[pyfunction] +#[pyo3( + signature = ( + key, + request_flags=None, + ), + text_signature = "(key: bytes, request_flags: Optional[RequestFlags])", +)] +pub fn build_meta_delete( + py: Python, + key: &[u8], + request_flags: Option<&RequestFlags>, +) -> PyResult> { + match impl_build_cmd(b"md", key, None, request_flags, false) { + Some(buf) => Ok(PyBytes::new(py, &buf).into()), + None => Err(pyo3::exceptions::PyValueError::new_err("Key is too long")), + } +} + +#[pyfunction] +#[pyo3( + signature = ( + key, + request_flags=None, + ), + text_signature = "(key: bytes, request_flags: Optional[RequestFlags])", +)] +pub fn build_meta_arithmetic( + py: Python, + key: &[u8], + request_flags: Option<&RequestFlags>, +) -> PyResult> { + match impl_build_cmd(b"ma", key, None, request_flags, false) { + Some(buf) => Ok(PyBytes::new(py, &buf).into()), + None => Err(pyo3::exceptions::PyValueError::new_err("Key is too long")), + } +} + +#[pymodule] +fn meta_memcache_socket(_py: Python, module: &PyModule) -> PyResult<()> { + module.add_class::()?; + module.add_class::()?; + module.add_function(wrap_pyfunction!(parse_header, module)?)?; + module.add_function(wrap_pyfunction!(build_cmd, module)?)?; + module.add_function(wrap_pyfunction!(build_meta_get, module)?)?; + module.add_function(wrap_pyfunction!(build_meta_set, module)?)?; + module.add_function(wrap_pyfunction!(build_meta_delete, module)?)?; + module.add_function(wrap_pyfunction!(build_meta_arithmetic, module)?)?; + module.add("RESPONSE_VALUE", RESPONSE_VALUE)?; + module.add("RESPONSE_SUCCESS", RESPONSE_SUCCESS)?; + module.add("RESPONSE_NOT_STORED", RESPONSE_NOT_STORED)?; + module.add("RESPONSE_CONFLICT", RESPONSE_CONFLICT)?; + module.add("RESPONSE_MISS", RESPONSE_MISS)?; + module.add("RESPONSE_NOOP", RESPONSE_NOOP)?; + module.add("SET_MODE_ADD", SET_MODE_ADD)?; + module.add("SET_MODE_APPEND", SET_MODE_APPEND)?; + module.add("SET_MODE_PREPEND", SET_MODE_PREPEND)?; + module.add("SET_MODE_REPLACE", SET_MODE_REPLACE)?; + module.add("SET_MODE_SET", SET_MODE_SET)?; + module.add("MA_MODE_INC", MA_MODE_INC)?; + module.add("MA_MODE_DEC", MA_MODE_DEC)?; + Ok(()) +} diff --git a/src/request_flags.rs b/src/request_flags.rs new file mode 100644 index 0000000..fba3e2b --- /dev/null +++ b/src/request_flags.rs @@ -0,0 +1,322 @@ +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +use crate::{MA_MODE_INC, SET_MODE_SET}; + +#[pyclass] +pub struct RequestFlags { + #[pyo3(get, set)] + no_reply: bool, + #[pyo3(get, set)] + return_client_flag: bool, + #[pyo3(get, set)] + return_cas_token: bool, + #[pyo3(get, set)] + return_value: bool, + #[pyo3(get, set)] + return_ttl: bool, + #[pyo3(get, set)] + return_size: bool, + #[pyo3(get, set)] + return_last_access: bool, + #[pyo3(get, set)] + return_fetched: bool, + #[pyo3(get, set)] + return_key: bool, + #[pyo3(get, set)] + no_update_lru: bool, + #[pyo3(get, set)] + mark_stale: bool, + #[pyo3(get, set)] + cache_ttl: Option, + #[pyo3(get, set)] + recache_ttl: Option, + #[pyo3(get, set)] + vivify_on_miss_ttl: Option, + #[pyo3(get, set)] + client_flag: Option, + #[pyo3(get, set)] + ma_initial_value: Option, + #[pyo3(get, set)] + ma_delta_value: Option, + #[pyo3(get, set)] + cas_token: Option, + #[pyo3(get, set)] + opaque: Option>, + #[pyo3(get, set)] + mode: Option, +} + +impl RequestFlags { + pub fn push_bytes(&self, buf: &mut Vec) { + if self.no_reply { + buf.push(b' '); + buf.push(b'q'); + } + if self.return_client_flag { + buf.push(b' '); + buf.push(b'f'); + } + if self.return_cas_token { + buf.push(b' '); + buf.push(b'c'); + } + if self.return_value { + buf.push(b' '); + buf.push(b'v'); + } + if self.return_ttl { + buf.push(b' '); + buf.push(b't'); + } + if self.return_size { + buf.push(b' '); + buf.push(b's'); + } + if self.return_last_access { + buf.push(b' '); + buf.push(b'l'); + } + if self.return_fetched { + buf.push(b' '); + buf.push(b'h'); + } + if self.return_key { + buf.push(b' '); + buf.push(b'k'); + } + if self.no_update_lru { + buf.push(b' '); + buf.push(b'u'); + } + if self.mark_stale { + buf.push(b' '); + buf.push(b'I'); + } + if let Some(v) = self.cache_ttl { + buf.push(b' '); + buf.push(b'T'); + buf.extend_from_slice(&v.to_string().as_bytes()); + } + if let Some(v) = self.recache_ttl { + buf.push(b' '); + buf.push(b'R'); + buf.extend_from_slice(&v.to_string().as_bytes()); + } + if let Some(v) = self.vivify_on_miss_ttl { + buf.push(b' '); + buf.push(b'N'); + buf.extend_from_slice(&v.to_string().as_bytes()); + } + if let Some(v) = self.client_flag { + buf.push(b' '); + buf.push(b'F'); + buf.extend_from_slice(&v.to_string().as_bytes()); + } + if let Some(v) = self.ma_initial_value { + buf.push(b' '); + buf.push(b'J'); + buf.extend_from_slice(&v.to_string().as_bytes()); + } + if let Some(v) = self.ma_delta_value { + buf.push(b' '); + buf.push(b'D'); + buf.extend_from_slice(&v.to_string().as_bytes()); + } + if let Some(v) = self.cas_token { + buf.push(b' '); + buf.push(b'C'); + buf.extend_from_slice(&v.to_string().as_bytes()); + } + if let Some(v) = &self.opaque { + buf.push(b' '); + buf.push(b'O'); + buf.extend_from_slice(&v); + } + if let Some(v) = self.mode { + if v != SET_MODE_SET && v != MA_MODE_INC { + // Set/inc are the default, no need to send them + buf.push(b' '); + buf.push(b'M'); + buf.push(v); + } + } + } +} + +#[pymethods] +impl RequestFlags { + #[new] + #[pyo3( + signature = ( + /, + *, + no_reply=false, + return_client_flag=false, + return_cas_token=false, + return_value=false, + return_ttl=false, + return_size=false, + return_last_access=false, + return_fetched=false, + return_key=false, + no_update_lru=false, + mark_stale=false, + cache_ttl=None, + recache_ttl=None, + vivify_on_miss_ttl=None, + client_flag=None, + ma_initial_value=None, + ma_delta_value=None, + cas_token=None, + opaque=None, + mode=None + ), + text_signature = "(*, + no_reply: bool = False, + return_client_flag: bool = False, + return_cas_token: bool = False, + return_value = False + return_ttl: bool = False, + return_size: bool = False, + return_last_access: bool = False, + return_fetched: bool = False, + return_key: bool = False, + no_update_lru: bool = False, + mark_stale: bool = False, + cache_ttl: Optional[int] = None, + recache_ttl: Optional[int] = None, + vivify_on_miss_ttl: Optional[int] = None, + client_flag: Optional[int] = None, + ma_initial_value: Optional[int] = None, + ma_delta_value: Optional[int] = None, + cas_token: Optional[int] = None, + opaque: Optional[bytes] = None, + mode: Optional[int] = None)" + )] + pub fn new( + no_reply: bool, + return_client_flag: bool, + return_cas_token: bool, + return_value: bool, + return_ttl: bool, + return_size: bool, + return_last_access: bool, + return_fetched: bool, + return_key: bool, + no_update_lru: bool, + mark_stale: bool, + cache_ttl: Option, + recache_ttl: Option, + vivify_on_miss_ttl: Option, + client_flag: Option, + ma_initial_value: Option, + ma_delta_value: Option, + cas_token: Option, + opaque: Option>, + mode: Option, + ) -> Self { + RequestFlags { + no_reply, + return_client_flag, + return_cas_token, + return_value, + return_ttl, + return_size, + return_last_access, + return_fetched, + return_key, + no_update_lru, + mark_stale, + cache_ttl, + recache_ttl, + vivify_on_miss_ttl, + client_flag, + ma_initial_value, + ma_delta_value, + cas_token, + opaque, + mode, + } + } + + pub fn copy(&self) -> Self { + RequestFlags { + no_reply: self.no_reply, + return_client_flag: self.return_client_flag, + return_cas_token: self.return_cas_token, + return_value: self.return_value, + return_ttl: self.return_ttl, + return_size: self.return_size, + return_last_access: self.return_last_access, + return_fetched: self.return_fetched, + return_key: self.return_key, + no_update_lru: self.no_update_lru, + mark_stale: self.mark_stale, + cache_ttl: self.cache_ttl, + recache_ttl: self.recache_ttl, + vivify_on_miss_ttl: self.vivify_on_miss_ttl, + client_flag: self.client_flag, + ma_initial_value: self.ma_initial_value, + ma_delta_value: self.ma_delta_value, + cas_token: self.cas_token, + opaque: self.opaque.clone(), + mode: self.mode, + } + } + + pub fn __eq__(&self, other: &Self) -> bool { + self.no_reply == other.no_reply + && self.return_client_flag == other.return_client_flag + && self.return_cas_token == other.return_cas_token + && self.return_value == other.return_value + && self.return_ttl == other.return_ttl + && self.return_size == other.return_size + && self.return_last_access == other.return_last_access + && self.return_fetched == other.return_fetched + && self.return_key == other.return_key + && self.no_update_lru == other.no_update_lru + && self.mark_stale == other.mark_stale + && self.cache_ttl == other.cache_ttl + && self.recache_ttl == other.recache_ttl + && self.vivify_on_miss_ttl == other.vivify_on_miss_ttl + && self.client_flag == other.client_flag + && self.ma_initial_value == other.ma_initial_value + && self.ma_delta_value == other.ma_delta_value + && self.cas_token == other.cas_token + && self.opaque == other.opaque + && self.mode == other.mode + } + + pub fn __str__(&self) -> String { + format!( + "RequestFlags(no_reply={:?}, return_client_flag={:?}, return_cas_token={:?}, return_value={:?}, return_ttl={:?}, return_size={:?}, return_last_access={:?}, return_fetched={:?}, return_key={:?}, no_update_lru={:?}, mark_stale={:?}, cache_ttl={:?}, recache_ttl={:?}, vivify_on_miss_ttl={:?}, client_flag={:?}, ma_initial_value={:?}, ma_delta_value={:?}, cas_token={:?}, opaque={:?}, mode={:?})", + self.no_reply, + self.return_client_flag, + self.return_cas_token, + self.return_value, + self.return_ttl, + self.return_size, + self.return_last_access, + self.return_fetched, + self.return_key, + self.no_update_lru, + self.mark_stale, + self.cache_ttl, + self.recache_ttl, + self.vivify_on_miss_ttl, + self.client_flag, + self.ma_initial_value, + self.ma_delta_value, + self.cas_token, + self.opaque, + self.mode, + ) + } + + pub fn to_bytes(&self, py: Python) -> Py { + let mut flags: Vec = Vec::new(); + self.push_bytes(&mut flags); + PyBytes::new(py, &flags).into() + } +} diff --git a/src/response_flags.rs b/src/response_flags.rs new file mode 100644 index 0000000..0496701 --- /dev/null +++ b/src/response_flags.rs @@ -0,0 +1,248 @@ +use atoi::FromRadix10Checked; +use pyo3::prelude::*; + +#[inline] +fn find_space_or_end(header: &[u8], start: usize) -> usize { + let mut n = start; + while n < header.len() { + if header[n] == b' ' { + break; + } + n += 1; + } + n +} + +#[inline] +fn get_u32_value(header: &[u8], start: usize) -> (Option, usize) { + match u32::from_radix_10_checked(&header[start..]) { + (Some(v), len) if len > 0 => (Some(v), start + len), + _ => (None, find_space_or_end(&header, start)), + } +} + +#[inline] +fn get_i32_value(header: &[u8], start: usize) -> (Option, usize) { + match i32::from_radix_10_checked(&header[start..]) { + (Some(v), len) if len > 0 => (Some(v), start + len), + _ => (None, find_space_or_end(&header, start)), + } +} + +#[pyclass] +pub struct ResponseFlags { + #[pyo3(get)] + pub cas_token: Option, + #[pyo3(get)] + pub fetched: Option, + #[pyo3(get)] + pub last_access: Option, + #[pyo3(get)] + pub ttl: Option, + #[pyo3(get)] + pub client_flag: Option, + #[pyo3(get)] + pub win: Option, + #[pyo3(get)] + pub stale: bool, + #[pyo3(get)] + pub size: Option, + #[pyo3(get)] + pub opaque: Option>, +} + +#[pymethods] +impl ResponseFlags { + #[new] + #[pyo3( + signature = ( + /, + *, + cas_token=None, + fetched=None, + last_access=None, + ttl=None, + client_flag=None, + win=None, + stale=false, + size=None, + opaque=None, + ), + text_signature = "(*, + cas_token=None, + fetched=None, + last_access=None, + ttl=None, + client_flag=None, + win=None, + stale=False, + size=None, + opaque=None)" + )] + fn new( + cas_token: Option, + fetched: Option, + last_access: Option, + ttl: Option, + client_flag: Option, + win: Option, + stale: Option, + size: Option, + opaque: Option>, + ) -> Self { + ResponseFlags { + cas_token, + fetched, + last_access, + ttl, + client_flag, + win, + stale: stale.unwrap_or(false), + size, + opaque, + } + } + + pub fn __eq__(&self, other: &Self) -> bool { + self.cas_token == other.cas_token + && self.fetched == other.fetched + && self.last_access == other.last_access + && self.ttl == other.ttl + && self.client_flag == other.client_flag + && self.win == other.win + && self.stale == other.stale + && self.size == other.size + && self.opaque == other.opaque + } + + pub fn __str__(&self) -> String { + format!( + "ResponseFlags(cas_token={:?}, fetched={:?}, last_access={:?}, ttl={:?}, client_flag={:?}, win={:?}, stale={}, size={:?}, opaque={:?})", + self.cas_token, + self.fetched, + self.last_access, + self.ttl, + self.client_flag, + self.win, + self.stale, + self.size, + self.opaque, + ) + } + + #[staticmethod] + pub fn from_success_header(header: &[u8]) -> Self { + return ResponseFlags::parse_flags(header, 3); + } + + #[staticmethod] + pub fn from_value_header(header: &[u8]) -> Option<(u32, Self)> { + let size_start: usize = 3; + if header.len() < size_start + 1 { + return None; + } + // let (size, pos) = u32::from_radix_10_checked(&header[size_start..]); + match u32::from_radix_10_checked(&header[size_start..]) { + (Some(size), pos) if pos > 0 => { + let flags = ResponseFlags::parse_flags(header, size_start + pos); + Some((size, flags)) + } + _ => None, + } + } + + #[staticmethod] + pub fn parse_flags(header: &[u8], start: usize) -> Self { + let mut cas_token: Option = None; + let mut fetched: Option = None; + let mut last_access: Option = None; + let mut ttl: Option = None; + let mut client_flag: Option = None; + let mut win: Option = None; + let mut stale: bool = false; + let mut size: Option = None; + let mut opaque: Option> = None; + + let mut n = start; + while n < header.len() { + let flag = header[n]; + // Move past the flag + n += 1; + match flag { + b' ' => { + // Space, skip it + continue; + } + b'c' => { + // cas_token flag (u32) + (cas_token, n) = get_u32_value(&header, n); + } + b'h' => { + // fetched flag (bool) encoded as 1 or 0 + fetched = match header[n] { + b'1' => Some(true), + b'0' => Some(false), + _ => None, + }; + n = find_space_or_end(&header, n + 1); + } + b'l' => { + // last_access flag (u32) + (last_access, n) = get_u32_value(&header, n); + } + b't' => { + // ttl flag (i32) encoded as -1 (for no ttl) or a positive number + if n < header.len() && header[n] == b'-' { + ttl = Some(-1); + n = find_space_or_end(&header, n) + } else { + (ttl, n) = get_i32_value(&header, n); + } + } + b'f' => { + // client_flag flag (u32) + (client_flag, n) = get_u32_value(&header, n); + } + b'W' => { + // win flag (bool), no value + win = Some(true); + } + b'Z' => { + // lost flag (bool), no value + win = Some(false); + } + b'X' => { + // stale flag (bool), no value + stale = true; + } + b's' => { + // size flag (u32) + (size, n) = get_u32_value(&header, n); + } + b'O' => { + // opaque flag (bytes) + let start = n; + n = find_space_or_end(&header, start); + opaque = Some(header[start..n].to_vec()); + } + _ => { + // Unknown flag, skip it + n = find_space_or_end(&header, n); + } + } + // n points now to a space, so continue past it + n += 1; + } + ResponseFlags { + cas_token, + fetched, + last_access, + ttl, + client_flag, + win, + stale, + size, + opaque, + } + } +}