gvsafronov/futriix-json
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Contributing valkeyJSON module (#1)

Browse Source 
Initial contribution for ValkeyJSON based off of Amazon implementation.
This commit is contained in:
Roshan Khatri 2024-11-29 07:47:54 -08:00 committed by GitHub
parent c1423d7e4b
commit 926b6fd6fe
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
58 changed files with 27928 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
.github/workflows/ci.yml vendored Normal file
View File

@ -0,0 +1,25 @@
name: CI
on: [push, pull_request]
jobs:
build-ubuntu-latest:
runs-on: ubuntu-latest
strategy:
fail-fast: false
matrix:
server_version: ['unstable', '8.0.0']
steps:
- uses: actions/checkout@v4
- name: Set the server verison for python integeration tests
run: echo "SERVER_VERSION=${{ matrix.server_version }}" >> $GITHUB_ENV
- name: Set up Python
uses: actions/setup-python@v3
with:
python-version: '3.9'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install -r requirements.txt
- name: Build and Run tests.
run: ./build.sh

31
.gitignore vendored Normal file
View File

@ -0,0 +1,31 @@
# IDE files
.idea/
*.vscode
.vscode/*
# Build temp files
*build
cmake-build-*/
# Auto-generated files
**/__pycache__/*
test-data
*.pyc
*.bin
*.o
*.xo
*.so
*.d
*.a
*.log
*.out
# Others
.DS_Store
.attach_pid*
venv/
core.*
valkeytests
**/include
**/report.html
**/assets

165
CMakeLists.txt Normal file
View File

@ -0,0 +1,165 @@
cmake_minimum_required(VERSION 3.17)
include(FetchContent)
include(ExternalProject)
# Detect the system architecture
EXECUTE_PROCESS(
COMMAND uname -m
COMMAND tr -d '\n'
OUTPUT_VARIABLE ARCHITECTURE
)
if("${ARCHITECTURE}" STREQUAL "x86_64")
message("Building JSON for x86_64")
elseif("${ARCHITECTURE}" STREQUAL "aarch64")
message("Building JSON for aarch64")
else()
message(FATAL_ERROR "Unsupported architecture: ${ARCHITECTURE}. JSON is only supported on x86_64 and aarch64.")
endif()
# Project definition
project(ValkeyJSONModule VERSION 1.0 LANGUAGES C CXX)
# Set the name of the JSON shared library
set(JSON_MODULE_LIB json)
# Define the Valkey directories
set(VALKEY_DOWNLOAD_DIR "${CMAKE_BINARY_DIR}/_deps/valkey-src")
set(VALKEY_BIN_DIR "${CMAKE_BINARY_DIR}/_deps/valkey-src/src/valkey/src")
# Download and build Valkey
ExternalProject_Add(
valkey
GIT_REPOSITORY https://github.com/valkey-io/valkey.git # Replace with actual URL
GIT_TAG ${VALKEY_VERSION}
PREFIX ${VALKEY_DOWNLOAD_DIR}
CONFIGURE_COMMAND ""
BUILD_COMMAND make distclean && make -j
INSTALL_COMMAND ""
BUILD_IN_SOURCE 1
)
# Define the paths for the copied files
set(VALKEY_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src/include")
set(VALKEY_BINARY_DEST "${CMAKE_CURRENT_SOURCE_DIR}/tst/integration/.build/binaries/${VALKEY_VERSION}")
ExternalProject_Add_Step(
valkey
copy_header_files
COMMENT "Copying header files to include/ directory"
DEPENDEES download
DEPENDERS configure
COMMAND ${CMAKE_COMMAND} -E make_directory ${VALKEY_INCLUDE_DIR}
COMMAND ${CMAKE_COMMAND} -E copy ${VALKEY_DOWNLOAD_DIR}/src/valkey/src/valkeymodule.h ${VALKEY_INCLUDE_DIR}/valkeymodule.h
ALWAYS 1
)
# Copy header and binary after Valkey make
add_custom_command(TARGET valkey
POST_BUILD
COMMAND ${CMAKE_COMMAND} -E make_directory ${VALKEY_BINARY_DEST}
COMMAND ${CMAKE_COMMAND} -E copy ${VALKEY_BIN_DIR}/valkey-server ${VALKEY_BINARY_DEST}/valkey-server
COMMENT "Copied valkeymodule.h and valkey-server to destination directories"
)
# Define valkey-bloom branch
set(VALKEY_BLOOM_BRANCH "unstable" CACHE STRING "Valkey-bloom branch to use")
# Set the download directory for Valkey-bloom
set(VALKEY_BLOOM_DOWNLOAD_DIR "${CMAKE_CURRENT_BINARY_DIR}/_deps/valkey-bloom-src")
# Download valkey-bloom
ExternalProject_Add(
valkey-bloom
GIT_REPOSITORY https://github.com/valkey-io/valkey-bloom.git
GIT_TAG ${VALKEY_BLOOM_BRANCH}
GIT_SHALLOW TRUE
PREFIX "${VALKEY_BLOOM_DOWNLOAD_DIR}"
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
INSTALL_COMMAND ""
)
# Step to copy pytest files
ExternalProject_Add_Step(
valkey-bloom
copy_pytest_files
COMMENT "Copying pytest files to tst/integration directory"
DEPENDEES build
COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_SOURCE_DIR}/tst/integration
COMMAND ${CMAKE_COMMAND} -E copy_directory ${VALKEY_BLOOM_DOWNLOAD_DIR}/src/valkey-bloom/tests/valkeytests ${CMAKE_CURRENT_SOURCE_DIR}/tst/integration/valkeytests
)
# Enable instrumentation options if requested
if("$ENV{INSTRUMENT_V2PATH}" STREQUAL "yes")
add_compile_definitions(INSTRUMENT_V2PATH)
message("Enabled INSTRUMENT_V2PATH")
endif()
# Disable Doxygen documentation generation
set(BUILD_DOCUMENTATION OFF)
# When CODE_COVERAGE is ON, the package is built twice, once for debug and once for release.
# To fix the problem, disable the code coverage.
set(CODE_COVERAGE OFF)
# Fix for linking error when code coverage is enabled on ARM
if(CODE_COVERAGE AND CMAKE_BUILD_TYPE STREQUAL "Debug")
add_link_options("--coverage")
endif()
# Set C & C++ standard versions
set(CMAKE_C_STANDARD 11)
set(CMAKE_C_STANDARD_REQUIRED True)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED True)
# Always include debug symbols and optimize the code
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -O3 -g -fno-omit-frame-pointer")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -g -fno-omit-frame-pointer")
# RapidJSON SIMD optimization
if("${ARCHITECTURE}" STREQUAL "x86_64")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=nehalem")
elseif("${ARCHITECTURE}" STREQUAL "aarch64")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=armv8-a")
else()
message(FATAL_ERROR "Unsupported architecture: ${ARCHITECTURE}. JSON is only supported on x86_64 and aarch64.")
endif()
# Additional flags for all architectures
set(ADDITIONAL_FLAGS "-fPIC")
# Compiler warning flags
set(C_WARNING "-Wall -Werror -Wextra")
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${ADDITIONAL_FLAGS} ${C_WARNING}")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${ADDITIONAL_FLAGS} ${C_WARNING}")
# Fetch RapidJSON
FetchContent_Declare(
rapidjson
GIT_REPOSITORY https://github.com/Tencent/rapidjson.git
GIT_TAG 0d4517f15a8d7167ba9ae67f3f22a559ca841e3b
)
# Disable RapidJSON tests and examples
set(RAPIDJSON_BUILD_TESTS OFF CACHE BOOL "Build rapidjson tests" FORCE)
set(RAPIDJSON_BUILD_EXAMPLES OFF CACHE BOOL "Build rapidjson examples" FORCE)
set(RAPIDJSON_BUILD_DOC OFF CACHE BOOL "Build rapidjson documentation" FORCE)
# Make Rapidjson available
FetchContent_MakeAvailable(rapidjson)
add_subdirectory(src)
add_subdirectory(tst)
add_custom_target(test
COMMENT "Run JSON integration tests."
USES_TERMINAL
COMMAND rm -rf ${CMAKE_BINARY_DIR}/tst/integration
COMMAND mkdir -p ${CMAKE_BINARY_DIR}/tst/integration
COMMAND cp -rp ${CMAKE_SOURCE_DIR}/tst/integration/. ${CMAKE_BINARY_DIR}/tst/integration/
COMMAND echo "[TARGET] begin integration tests"
COMMAND ${CMAKE_SOURCE_DIR}/tst/integration/run.sh "test" ${CMAKE_SOURCE_DIR}
COMMAND echo "[TARGET] end integration tests")

86
README.md Normal file
View File

@ -0,0 +1,86 @@
# ValkeyJSON
ValkeyJSON is a C++ Valkey-Module that provides native JSON (JavaScript Object Notation) support for Valkey. The implementation complies with RFC7159 and ECMA-404 JSON data interchange standards. Users can natively store, query, and modify JSON data structures using the JSONPath query language. The query expressions support advanced capabilities including wildcard selections, filter expressions, array slices, union operations, and recursive searches.
ValkeyJSON leverages [RapidJSON](https://rapidjson.org/), a high-performance JSON parser and generator for C++, chosen for its small footprint and exceptional performance and memory efficiency. As a header-only library with no external dependencies, RapidJSON provides robust Unicode support while maintaining a compact memory profile of just 16 bytes per JSON value on most 32/64-bit machines.
## Building and Testing
#### To build the module and run tests
```text
# Builds the valkey-server (unstable) for integration testing.
export SERVER_VERSION=unstable
./build.sh
# Builds the valkey-server (8.0.0) for integration testing.
export SERVER_VERSION=8.0.0
./build.sh
```
#### To build just the module
```text
mdkir build
cd build
cmake .. -DVALKEY_VERSION=unstable
make
```
#### To run all unit tests:
```text
cd build
make -j unit
```
#### To run all integration tests:
```text
make -j test
```
## Load the Module
To test the module with a Valkey, you can load the module using any of the following ways:
#### Using valkey.conf:
```
1. Add the following to valkey.conf:
loadmodule /path/to/libjson.so
2. Start valkey-server:
valkey-server /path/to/valkey.conf
```
#### Starting valkey with --loadmodule option:
```text
valkey-server --loadmodule /path/to/libjson.so
```
#### Using Valkey command MODULE LOAD:
```
1. Connect to a running Valkey instance using valkey-cli
2. Execute Valkey command:
MODULE LOAD /path/to/libjson.so
```
## Supported Module Commands
```text
JSON.ARRAPPEND
JSON.ARRINDEX
JSON.ARRINSERT
JSON.ARRLEN
JSON.ARRPOP
JSON.ARRTRIM
JSON.CLEAR
JSON.DEBUG
JSON.DEL
JSON.FORGET
JSON.GET
JSON.MGET
JSON.MSET
JSON.NUMINCRBY
JSON.NUMMULTBY
JSON.OBJLEN
JSON.OBJKEYS
JSON.RESP
JSON.SET
JSON.STRAPPEND
JSON.STRLEN
JSON.TOGGLE
JSON.TYPE
```

63
build.sh Executable file
View File

@ -0,0 +1,63 @@
#!/bin/bash
# Script to build valkeyJSON module, build it and generate .so files, run unit and integration tests.
# # Exit the script if any command fails
set -e
SCRIPT_DIR=$(pwd)
echo "Script Directory: $SCRIPT_DIR"
# Ensure SERVER_VERSION environment variable is set
if [ -z "$SERVER_VERSION" ]; then
echo "WARNING: SERVER_VERSION environment variable is not set. Defaulting to unstable."
export SERVER_VERSION="unstable"
fi
if [ "$SERVER_VERSION" != "unstable" ] && [ "$SERVER_VERSION" != "8.0.0" ] ; then
echo "ERROR: Unsupported version - $SERVER_VERSION"
exit 1
fi
# Variables
BUILD_DIR="$SCRIPT_DIR/build"
# Build the Valkey JSON module using CMake
echo "Building ValkeyJSON module..."
if [ ! -d "$BUILD_DIR" ]; then
mkdir $BUILD_DIR
fi
cd $BUILD_DIR
cmake .. -DVALKEY_VERSION=$SERVER_VERSION
make
# Running the Valkey JSON unit tests.
echo "Running Unit Tests..."
make -j unit
cd $SCRIPT_DIR
REQUIREMENTS_FILE="requirements.txt"
# Check if pip is available
if command -v pip > /dev/null 2>&1; then
echo "Using pip to install packages..."
pip install -r "$SCRIPT_DIR/$REQUIREMENTS_FILE"
# Check if pip3 is available
elif command -v pip3 > /dev/null 2>&1; then
echo "Using pip3 to install packages..."
pip3 install -r "$SCRIPT_DIR/$REQUIREMENTS_FILE"
else
echo "Error: Neither pip nor pip3 is available. Please install Python package installer."
exit 1
fi
export MODULE_PATH="$SCRIPT_DIR/build/src/libjson.so"
# Running the Valkey JSON integration tests.
echo "Running the integration tests..."
cd $BUILD_DIR
make -j test
echo "Build and Integration Tests succeeded"

3
requirements.txt Normal file
View File

@ -0,0 +1,3 @@
valkey
pytest==4
pytest-html

40
src/CMakeLists.txt Normal file
View File

@ -0,0 +1,40 @@
message("src/CMakeLists.txt: Build JSON")
set(OBJECT_TARGET json-objects CACHE INTERNAL "Object target for json module")
add_library(${OBJECT_TARGET} OBJECT "")
# Build with C11 & C++17
set_target_properties(
${OBJECT_TARGET}
PROPERTIES
C_STANDARD 11
C_STANDARD_REQUIRED ON
CXX_STANDARD 17
CXX_STANDARD_REQUIRED ON
POSITION_INDEPENDENT_CODE ON
)
target_include_directories(${OBJECT_TARGET}
# Need to make the source files public within CMake
# so that they are used when building the tests.
PUBLIC
${CMAKE_CURRENT_SOURCE_DIR}
${rapidjson_SOURCE_DIR}/include
)
# Add source files for the JSON module
target_sources(${OBJECT_TARGET}
PRIVATE
json/json.cc
json/dom.cc
json/alloc.cc
json/util.cc
json/stats.cc
json/selector.cc
json/keytable.cc
json/memory.cc
json/json_api.cc
)
add_library(${JSON_MODULE_LIB} SHARED $<TARGET_OBJECTS:${OBJECT_TARGET}>)

2
src/json/CPPLINT.cfg Normal file
View File

@ -0,0 +1,2 @@
filter=-build/include_order,-legal/copyright,-whitespace/braces,-build/c++11,-runtime/references,-build/include_what_you_use,-readability/casting,-build/header_guard,-runtime/int,-build/namespaces,-runtime/explicit
linelength=120

59
src/json/alloc.cc Normal file
View File

@ -0,0 +1,59 @@
#include "json/memory.h"
#include "json/alloc.h"
#include "json/stats.h"
#include <cstdio>
#include <cstring>
extern "C" {
#define VALKEYMODULE_EXPERIMENTAL_API
#include <./include/valkeymodule.h>
}
void *dom_alloc(size_t size) {
void *ptr = memory_alloc(size);
// actually allocated size may not be same as the requested size
size_t real_size = memory_allocsize(ptr);
jsonstats_increment_used_mem(real_size);
return ptr;
}
void dom_free(void *ptr) {
size_t size = memory_allocsize(ptr);
memory_free(ptr);
jsonstats_decrement_used_mem(size);
}
void *dom_realloc(void *orig_ptr, size_t new_size) {
// We need to handle the following two edge cases first. Otherwise, the following
// calculation of the incremented/decremented amount will fail.
if (new_size == 0 && orig_ptr != nullptr) {
dom_free(orig_ptr);
return nullptr;
}
if (orig_ptr == nullptr) return dom_alloc(new_size);
size_t orig_size = memory_allocsize(orig_ptr);
void *new_ptr = memory_realloc(orig_ptr, new_size);
// actually allocated size may not be same as the requested size
size_t real_new_size = memory_allocsize(new_ptr);
if (real_new_size > orig_size)
jsonstats_increment_used_mem(real_new_size - orig_size);
else if (real_new_size < orig_size)
jsonstats_decrement_used_mem(orig_size - real_new_size);
return new_ptr;
}
char *dom_strdup(const char *s) {
size_t size = strlen(s) + 1;
char *dup = static_cast<char*>(dom_alloc(size));
strncpy(dup, s, size);
return dup;
}
char *dom_strndup(const char *s, const size_t n) {
char *dup = static_cast<char*>(dom_alloc(n + 1));
strncpy(dup, s, n);
dup[n] = '\0';
return dup;
}

29
src/json/alloc.h Normal file
View File

@ -0,0 +1,29 @@
/**
* This C module is the JSON memory allocator (also called DOM allocator), which wraps around Valkey's built-in
* allocation functions - ValkeyModule_Alloc, ValkeyModule_Free and ValkeyModule_Realloc. All memory allocations,
* permanent or transient, should be done through this interface, so that allocated memories are correctly
* reported to the Valkey engine (MEMORY STATS).
*
* Besides correctly reporting memory usage to Valkey, it also provides a facility to track memory usage of JSON
* objects, so that we can achieve the following:
* 1. To track total memory allocated to JSON objects. This is done through an atomic global counter. Note that
* Valkey engine only reports total memories for all keys, not by key type. This JSON memory allocator overcomes
* such deficiency.
* 2. To track each JSON document object's memory size. This is done through a thread local counter. With the ability
* to track individual document's footprint, we can maintain a few interesting histograms that will provide
* insights into data distribution and API access patterns.
*/
#ifndef VALKEYJSONMODULE_ALLOC_H_
#define VALKEYJSONMODULE_ALLOC_H_
#include <stddef.h>
#include "json/memory.h"
void *dom_alloc(size_t size);
void dom_free(void *ptr);
void *dom_realloc(void *orig_ptr, size_t new_size);
char *dom_strdup(const char *s);
char *dom_strndup(const char *s, const size_t n);
#endif // VALKEYJSONMODULE_ALLOC_H_

1624
src/json/dom.cc Normal file
View File

File diff suppressed because it is too large Load Diff

545
src/json/dom.h Normal file
View File

@ -0,0 +1,545 @@
/**
* DOM (Document Object Model) interface for JSON.
* The DOM module provides the following functions:
* 1. Parsing and validating an input JSON string buffer
* 2. Deserializing a JSON string into document object
* 2. Serializing a document object into JSON string
* 3. JSON CRUD operations: search, insert, update and delete
*
* Design Considerations:
* 1. Memory management: All memory management must be handled by the JSON allocator.
* - For memories allocated by our own code:
* All allocations and de-allocations must be done through:
* dom_alloc, dom_free, dom_realloc, dom_strdup, and dom_strndup
* - For objects allocated by RapidJSON library:
* Our solution is to use a custom memory allocator class as template, so as to instruct RapidJSON to the JSON
* allocator. The custom allocator works under the hood and is not exposed through this interface.
* 2. If a method returns to the caller a heap-allocated object, it must be documented.
* The caller is responsible for releasing the memory after consuming it.
* 3. Generally speaking, interface methods should not have Valkey module types such as ValkeyModuleCtx or
* ValkeyModuleString, because that would make unit tests hard to write, unless gmock classes have been developed.
*
* Coding Conventions & Best Practices:
* 1. Error handling: If a method may fail, the return type should be enum JsonUtilCode.
* 2. Output parameters: Output parameters should be placed at the end. Output parameters should be initialized at the
* beginning of the method. It should not require the caller to do any initialization before invoking the method.
* 3. Every public interface method declared in this file should be prefixed with "dom_".
*/
#ifndef VALKEYJSONMODULE_JSON_DOM_H_
#define VALKEYJSONMODULE_JSON_DOM_H_
#include <stdlib.h>
#include <string>
#include "json/util.h"
#include "json/alloc.h"
#include "json/rapidjson_includes.h"
class ReplyBuffer : public rapidjson::StringBuffer {
public:
ReplyBuffer(ValkeyModuleCtx *_ctx, bool) : rapidjson::StringBuffer(), ctx(_ctx) {}
ReplyBuffer() : rapidjson::StringBuffer(), ctx(nullptr) {}
void Initialize(ValkeyModuleCtx *_ctx, bool) { ctx = _ctx; }
void Reply() { ValkeyModule_ReplyWithStringBuffer(ctx, GetString(), GetLength()); }
private:
ValkeyModuleCtx *ctx;
};
extern "C" {
#define VALKEYMODULE_EXPERIMENTAL_API
#include <./include/valkeymodule.h>
}
/**
* This is a custom allocator for RapidJSON. It delegates memory management to the JSON allocator, so that
* memory allocated by the underlying RapidJSON library can be correctly reported to Valkey engine. The class
* is passed into rapidjson::GenericDocument and rapidjson::GenericValue as template, which is the way to tell
* RapidJSON to use a custom allocator.
*/
class RapidJsonAllocator {
public:
RapidJsonAllocator();
void *Malloc(size_t size) {
return dom_alloc(size);
}
void *Realloc(void *originalPtr, size_t /*originalSize*/, size_t newSize) {
return dom_realloc(originalPtr, newSize);
}
static void Free(void *ptr) RAPIDJSON_NOEXCEPT {
dom_free(ptr);
}
bool operator==(const RapidJsonAllocator&) const RAPIDJSON_NOEXCEPT {
return true;
}
bool operator!=(const RapidJsonAllocator&) const RAPIDJSON_NOEXCEPT {
return false;
}
static const bool kNeedFree = true;
};
/**
* Now, wrap the RapidJSON objects (RJxxxxx) with our own objects (Jxxxxx). We wrap them to hide
* various RapidJSON oddities, simplify the syntax and to more explicitly express the semantics.
* For example, the details of allocators are largely hidden in the wrapped objects.
*
* We use four objects for all of our work. All of these objects descend from one of the three
* basic RapidJSON object types.
*
* RJValue (JValue): A JSON value. This is implemented as a node of a tree. This object doesn't
* differentiate between being the root of a tree or the root of a sub-tree.
* The full suite of RapidJSON value manipulation functions is available. Many
* of the RapidJSON value functions require an allocator. You must use the
* global "allocator".
*
* RJParser (JParser): This object contains a JValue into which you can deserialize a stream via the
* Parse/ParseStream member functions. The JValue created by the parsing routines
* is allocated using the dom_alloc/dom_free accounting. Typically, a JParser
* object is created on the run-time stack, filled with some serialized data
* and the then created JValue is moved into a destination location.
*
* JDocument This is the only object visible external to the dom layer. Externally, a
* JDocument is the Valkey data type for this module, i.e., the Valkey dictionary
* contains this pointer. Externally, this is an opaque data structure.
* Internally, it's implemented as a JValue plus a size and bucket number. The
* size is maintained as the memory size of the entire tree of JValues
* contained by the JDocument.
*/
typedef rapidjson::GenericValue<rapidjson::UTF8<>, RapidJsonAllocator> RJValue;
// A JValue is an RJValue without any local augmentation of change.
typedef RJValue JValue;
extern RapidJsonAllocator allocator;
/**
* A JDocument privately inherits from JValue. You must use the GetJValue() member
* to access the underlying JValue. This improves readability at the usage point.
*/
struct JDocument : JValue {
JDocument() : JValue(), size(0), bucket_id(0) {}
JValue& GetJValue() { return *this; }
const JValue& GetJValue() const { return *this; }
void SetJValue(JValue& rhs) { *static_cast<JValue *>(this) = rhs; }
size_t size:56; // Size of this document, maintained by the JSON layer, not here.
size_t bucket_id:8; // document histogram's bucket id. maintained by JSON layer, not here
void *operator new(size_t size) { return dom_alloc(size); }
void operator delete(void *ptr) { return dom_free(ptr); }
private:
//
// Since JDocument objects are 1:1 with Valkey Keys, you can't ever have an array of them.
//
void *operator new[](size_t); // Not defined anywhere, causes link error if used
void operator delete[](void *); // Not defined anywhere, causes link error if used
};
//
// typedef the RapidJSON objects we care about, name them RJxxxxx for clarity
//
typedef rapidjson::GenericDocument<rapidjson::UTF8<>, RapidJsonAllocator> RJParser;
/**
* A JParser privately inherits from RJParser, which inherits from RJValue. You must use the
* GetJValue() member to access the post Parse value.
*
*/
struct JParser : RJParser {
JParser() : RJParser(&allocator), allocated_size(0) {}
//
// Make these inner routines publicly visible
//
using RJParser::ParseStream;
using RJParser::HasParseError;
using RJParser::GetMaxDepth;
// Access the contained JValue
JValue& GetJValue() { return *this; }
//
// Translate rapidJSON parse error code into JsonUtilCode.
//
JsonUtilCode GetParseErrorCode() {
switch (GetParseError()) {
case rapidjson::kParseErrorTermination:
return JSONUTIL_DOCUMENT_PATH_LIMIT_EXCEEDED;
case rapidjson::kParseErrorNone:
ValkeyModule_Assert(false);
/* Fall Through, but not really */
default:
return JSONUTIL_JSON_PARSE_ERROR;
}
}
//
// When we parse an incoming string, we want to know how much member this will consume.
// So track it and retain it.
//
JParser& Parse(const char *json, size_t len);
JParser& Parse(const std::string_view &sv);
//
// This object holds a JValue which is the root of the parsed tree. The dom_alloc/dom_free
// machinery will track all memory allocations outside of this object, but the root JValue
// won't be covered by that. Because the JParser objects aren't created via new, they are
// created as stack variables. So we manually add that in, since it'll be charged to the
// destination when we actually move the value out.
//
size_t GetJValueSize() const { return allocated_size + sizeof(RJValue); }
private:
size_t allocated_size;
};
/* Parse input JSON string, validate syntax, and return a document object.
* This method can handle an input string that is not NULL terminated. One use case is that
* we call ValkeyModule_LoadStringBuffer() to load JSON data from RDB, which returns a string that
* is not automatically NULL terminated.
* Also note that if the input string has NULL character '\0' in the middle, the string
* terminates at the NULL character.
*
* @param json_buf - pointer to binary string buffer, which may not be NULL terminated.
* @param buf_len - length of the input string buffer, which may not be NULL terminated.
* @param doc - OUTPUT param, pointer to document pointer. The caller is responsible for calling
* dom_free_doc(JDocument*) to free the memory after it's consumed.
* @return JSONUTIL_SUCCESS for success, other code for failure.
*/
JsonUtilCode dom_parse(ValkeyModuleCtx *ctx, const char *json_buf, const size_t buf_len, JDocument **doc);
/* Free a document object */
void dom_free_doc(JDocument *doc);
/* Get document size */
size_t dom_get_doc_size(const JDocument *doc);
/* Set document size */
void dom_set_doc_size(JDocument *doc, const size_t size);
/* Get the document histogram's bucket ID */
size_t dom_get_bucket_id(const JDocument *doc);
/* Set the document histogram's bucket ID */
void dom_set_bucket_id(JDocument *doc, const uint32_t bucket_id);
/**
* Serialize a document into the given string stream.
* @param format - controls format of returned JSON string.
* if NULL, return JSON in compact format (no space, no indent, no newline).
* @param oss - output stream
*/
void dom_serialize(JDocument *doc, const PrintFormat *format, rapidjson::StringBuffer &oss);
/**
* Serialize a value into the given string stream.
* @param format - controls format of returned JSON string.
* if NULL, return JSON in compact format (no space, no indent, no newline).
* @param oss - output stream
* @param json_len - OUTPUT param, *json_len is length of JSON string.
*/
void dom_serialize_value(const JValue &val, const PrintFormat *format, rapidjson::StringBuffer &oss);
/**
* Get the root value of the document.
*/
JValue& dom_get_value(JDocument &doc);
/* Set value at the path.
* @param json_path: path that is compliant to the JSON Path syntax.
* @param is_create_only - indicates to create a new value.
* @param is_update_only - indicates to update an existing value.
* @return JSONUTIL_SUCCESS for success, other code for failure.
*/
JsonUtilCode dom_set_value(ValkeyModuleCtx *ctx, JDocument *doc, const char *json_path, const char *new_val_json,
size_t new_val_len, const bool is_create_only = false, const bool is_update_only = false);
inline JsonUtilCode dom_set_value(ValkeyModuleCtx *ctx, JDocument *doc, const char *json_path, const char *new_val_json,
const bool is_create_only = false, const bool is_update_only = false) {
return dom_set_value(ctx, doc, json_path, new_val_json, strlen(new_val_json), is_create_only, is_update_only);
}
/* Get JSON value at the path.
* If the path is invalid, the method will return error code JSONUTIL_INVALID_JSON_PATH.
* If the path does not exist, the method will return error code JSONUTIL_JSON_PATH_NOT_EXIST.
*
* @param format - controls format of returned JSON string.
* if NULL, return JSON in compact format (no space, no indent, no newline).
* @param oss - output stream
* @return JSONUTIL_SUCCESS for success, other code for failure.
*/
template<typename T>
JsonUtilCode dom_get_value_as_str(JDocument *doc, const char *json_path, const PrintFormat *format,
T &oss, const bool update_stats = true);
/* Get JSON values at multiple paths. Values at multiple paths will be aggregated into a JSON object,
* in which each path is a key.
* If the path is invalid, the method will return error code JSONUTIL_INVALID_JSON_PATH.
* If the path does not exist, the method will return error code JSONUTIL_JSON_PATH_NOT_EXIST.
*
* @param format - controls format of returned JSON string.
* if NULL, return JSON in compact format (no space, no indent, no newline).
* @param oss - output stream, the string represents an aggregated JSON object in which each path is a key.
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_get_values_as_str(JDocument *doc, const char **paths, const int num_paths,
PrintFormat *format, ReplyBuffer &oss, const bool update_stats = true);
/**
* Delete JSON values at the given path.
* @num_vals_deleted number of values deleted
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_delete_value(JDocument *doc, const char *json_path, size_t &num_vals_deleted);
/* Increment the JSON value by a given number.
* @param out_val OUTPUT parameter, pointer to new value
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_increment_by(JDocument *doc, const char *json_path, const JValue *incr_by,
jsn::vector<double> &out_vals, bool &is_v2_path);
/* Multiply the JSON value by a given number.
* @param out_val OUTPUT parameter, pointer to new value
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_multiply_by(JDocument *doc, const char *json_path, const JValue *mult_by,
jsn::vector<double> &out_vals, bool &is_v2_path);
/* Toggle a JSON boolean between true and false.
* @param vec OUTPUT parameter, a vector of integers. 0: false, 1: true, -1: N/A - the source value is not boolean.
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_toggle(JDocument *doc, const char *path, jsn::vector<int> &vec, bool &is_v2_path);
/* Get the length of a JSON string value.
* @param vec OUTPUT parameter, a vector of string lengths
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_string_length(JDocument *doc, const char *path, jsn::vector<size_t> &vec, bool &is_v2_path);
/* Append a string to an existing JSON string value.
* @param vec OUTPUT parameter, a vector of new string lengths
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_string_append(JDocument *doc, const char *path, const char *json, const size_t json_len,
jsn::vector<size_t> &vec, bool &is_v2_path);
/**
* Get number of keys in the object at the given path.
* @param vec, OUTPUT parameter, a vector of object lengths
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_object_length(JDocument *doc, const char *path, jsn::vector<size_t> &vec, bool &is_v2_path);
/**
* Get keys in the object at the given path.
* @param vec OUTPUT parameter, a vector of vector of strings. In the first level vector, number of items is
* number of objects. In the second level vector, number of items is number keys in the object.
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_object_keys(JDocument *doc, const char *path,
jsn::vector<jsn::vector<jsn::string>> &vec, bool &is_v2_path);
/**
* Get number of elements in the array at the given path.
* @param vec OUTPUT parameter, a vector of array lengths
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_array_length(JDocument *doc, const char *path, jsn::vector<size_t> &vec, bool &is_v2_path);
/**
* Append a list of values to the array at the given path.
* @param vec OUTPUT parameter, a vector of new array lengths
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_array_append(ValkeyModuleCtx *ctx, JDocument *doc, const char *path,
const char **jsons, size_t *json_lens, const size_t num_values,
jsn::vector<size_t> &vec, bool &is_v2_path);
/**
* Remove and return element from the index in the array.
* Out of range index is rounded to respective array boundaries.
*
* @param index - position in the array to start popping from, defaults -1 , which means the last element.
* Negative value means position from the last element.
* @param vec - OUTPUT parameter, a vector of string streams, each containing JSON string of the popped element
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_array_pop(JDocument *doc, const char *path, int64_t index,
jsn::vector<rapidjson::StringBuffer> &vec, bool &is_v2_path);
/**
* Insert one or more json values into the array at path before the index.
* Inserting at index 0 prepends to the array.
* A negative index values in interpreted as starting from the end.
* The index must be in the array's range.
*
* @param vec OUTPUT parameter, a vector of new array lengths
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_array_insert(ValkeyModuleCtx *ctx, JDocument *doc, const char *path, int64_t index,
const char **jsons, size_t *json_lens, const size_t num_values,
jsn::vector<size_t> &vec, bool &is_v2_path);
/**
* Clear all the elements in an array or object.
* Return number of containers cleared.
*
* @param elements_cleared, OUTPUT parameter, number of elements cleared
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_clear(JDocument *doc, const char *path, size_t &elements_cleared);
/*
* Trim an array so that it becomes subarray [start, end], both inclusive.
* If the array is empty, do nothing, return 0.
* If start < 0, set it to 0.
* If stop >= size, set it to size-1
* If start >= size or start > stop, empty the array and return 0.
*
* @param start - start index, inclusive
* @param stop - stop index, inclusive
* @param vec, OUTPUT parameter, a vector of new array lengths
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_array_trim(JDocument *doc, const char *path, int64_t start, int64_t stop,
jsn::vector<size_t> &vec, bool &is_v2_path);
/**
* Search for the first occurrence of a scalar JSON value in an array.
* Out of range errors are treated by rounding the index to the array's start and end.
* If start > stop, return -1 (not found).
*
* @param scalar_val - scalar value to search for
* @param start - start index, inclusive
* @param stop - stop index, exclusive. 0 or -1 means the last element is included.
* @param vec OUTPUT parameter, a vector of matching indexes. -1 means value not found.
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_array_index_of(JDocument *doc, const char *path, const char *scalar_val,
const size_t scalar_val_len, int64_t start, int64_t stop,
jsn::vector<int64_t> &vec, bool &is_v2_path);
/* Get type of a JSON value.
* @param vec, OUTPUT parameter, a vector of value types.
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_value_type(JDocument *doc, const char *path, jsn::vector<jsn::string> &vec, bool &is_v2_path);
/*
* Return a JSON value in Valkey Serialization Protocol (RESP).
* If the value is container, the response is RESP array or nested array.
*
* JSON null is mapped to the RESP Null Bulk String.
* JSON boolean values are mapped to the respective RESP Simple Strings.
* JSON integer numbers are mapped to RESP Integers.
* JSON float or double numbers are mapped to RESP Bulk Strings.
* JSON Strings are mapped to RESP Bulk Strings.
* JSON Arrays are represented as RESP Arrays, where the first element is the simple string [,
* followed by the array's elements.
* JSON Objects are represented as RESP Arrays, where the first element is the simple string {,
* followed by key-value pairs, each of which is a RESP bulk string.
*/
JsonUtilCode dom_reply_with_resp(ValkeyModuleCtx *ctx, JDocument *doc, const char *path);
/* Get memory size of a JSON value.
* @param vec, OUTPUT parameter, vector of memory size.
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_mem_size(JDocument *doc, const char *path, jsn::vector<size_t> &vec, bool &is_v2_path,
bool default_path);
/* Get number of fields in a JSON value.
* @param vec, OUTPUT parameter, vector of number of fields.
* @return JSONUTIL_SUCCESS if success. Other codes indicate failure.
*/
JsonUtilCode dom_num_fields(JDocument *doc, const char *path, jsn::vector<size_t> &vec, bool &is_v2_path);
/**
* Get max path depth of a document.
*/
void dom_path_depth(JDocument *doc, size_t *depth);
/* Duplicate a JSON value. */
JDocument *dom_copy(const JDocument *source);
/*
* The dom_save and dom_load support the ability to save and load a single JSON document
* as a sequence of chunks of data. The advantage of chunking is that you never need a single
* buffer that's the size of the entire serialized object. Without chunking for a large JSON
* object you would need to reserve sufficient memory to serialize it en masse.
*
* dom_save synchronously serializes the object into a sequence of chunks, each chunk is
* delivered to a callback function for disposal.
*
* dom_load synchronously deserializes a series of chunks of data into a JSON object.
* It calls a callback which returns a chunk of data. That data is "owned" by dom_load until
* it passes the ownership of that chunk back to the caller via a callback. End of input can
* be signalled by returning a nullptr or a 0-length chunk of data.
*
* The usage of callbacks insulates the dom layer from any knowledge of the RDB format. That's
* exclusively done by the callback functions.
*/
/*
* save a document into rdb format.
* @param source JSON document to be saved
* @param rdb rdb file context
*/
void dom_save(const JDocument *source, ValkeyModuleIO *rdb, int encver);
/*
* load a document from rdb format
* @param dest output document pointer
* @param rdb rdb file context
* @param encver encoding version
*/
JsonUtilCode dom_load(JDocument **dest, ValkeyModuleIO *rdb, int encver);
/*
* Implement DEBUG DIGEST
*/
void dom_compute_digest(ValkeyModuleDigest *ctx, const JDocument *doc);
void dom_dump_value(JValue &v);
// Unit test
jsn::string validate(const JDocument *);
//
// JSON Validation functions
//
//
// Validates that all pointers contained within this JValue are valid.
//
// true => All good.
// false => Not all good.
//
bool ValidateJValue(JValue &v);
//
// This function dumps a JValue to an output stream (like an ostringstream)
//
// The structure of the object is dumped but no actual customer data is dumped.
// It's totally legal to call this function on a corrupted JValue, it'll avoid the bad mallocs
//
// Typical usage:
//
// std::ostringstream os;
// DumpRedactedJValue(os, v); // Don't specify level or index parameters, let them default
// ValkeyModule_Log(...., os.str());
//
void DumpRedactedJValue(std::ostream& os, const JValue &v, size_t level = 0, int index = -1);
//
// Same as above, except targets the Valkey Log
//
void DumpRedactedJValue(const JValue &v, ValkeyModuleCtx *ctx = nullptr, const char *level = "debug");
#endif // VALKEYJSONMODULE_JSON_DOM_H_

3231
src/json/json.cc Normal file
View File

File diff suppressed because it is too large Load Diff

22
src/json/json.h Normal file
View File

@ -0,0 +1,22 @@
#ifndef VALKEYJSONMODULE_JSON_H_
#define VALKEYJSONMODULE_JSON_H_
#include <stddef.h>
size_t json_get_max_document_size();
size_t json_get_defrag_threshold();
size_t json_get_max_path_limit();
size_t json_get_max_parser_recursion_depth();
size_t json_get_max_recursive_descent_tokens();
size_t json_get_max_query_string_size();
bool json_is_instrument_enabled_insert();
bool json_is_instrument_enabled_update();
bool json_is_instrument_enabled_delete();
bool json_is_instrument_enabled_dump_doc_before();
bool json_is_instrument_enabled_dump_doc_after();
bool json_is_instrument_enabled_dump_value_before_delete();
#define DOUBLE_CHARS_CUTOFF 24
#endif // VALKEYJSONMODULE_JSON_H_

120
src/json/json_api.cc Normal file
View File

@ -0,0 +1,120 @@
#include <cstddef>
#include <cstdlib>
#include <iostream>
#include <fstream>
#include <string>
#include "json/json_api.h"
#include "json/dom.h"
#include "json/memory.h"
extern ValkeyModuleType* DocumentType;
int is_json_key(ValkeyModuleCtx *ctx, ValkeyModuleKey *key) {
VALKEYMODULE_NOT_USED(ctx);
if (key == nullptr || ValkeyModule_KeyType(key) == VALKEYMODULE_KEYTYPE_EMPTY) return 0;
return (ValkeyModule_ModuleTypeGetType(key) == DocumentType? 1: 0);
}
int is_json_key2(ValkeyModuleCtx *ctx, ValkeyModuleString *keystr) {
ValkeyModuleKey *key = static_cast<ValkeyModuleKey*>(ValkeyModule_OpenKey(ctx, keystr, VALKEYMODULE_READ));
int is_json = is_json_key(ctx, key);
ValkeyModule_CloseKey(key);
return is_json;
}
static JDocument* get_json_document(ValkeyModuleCtx *ctx, const char *keyname, const size_t key_len) {
ValkeyModuleString *keystr = ValkeyModule_CreateString(ctx, keyname, key_len);
ValkeyModuleKey *key = static_cast<ValkeyModuleKey*>(ValkeyModule_OpenKey(ctx, keystr, VALKEYMODULE_READ));
if (!is_json_key(ctx, key)) {
ValkeyModule_CloseKey(key);
ValkeyModule_FreeString(ctx, keystr);
return nullptr;
}
JDocument *doc = static_cast<JDocument*>(ValkeyModule_ModuleTypeGetValue(key));
ValkeyModule_CloseKey(key);
ValkeyModule_FreeString(ctx, keystr);
return doc;
}
int get_json_value_type(ValkeyModuleCtx *ctx, const char *keyname, const size_t key_len, const char *path,
char **type, size_t *len) {
*type = nullptr;
*len = 0;
JDocument *doc = get_json_document(ctx, keyname, key_len);
if (doc == nullptr) return -1;
jsn::vector<jsn::string> vec;
bool is_v2_path;
JsonUtilCode rc = dom_value_type(doc, path, vec, is_v2_path);
if (rc != JSONUTIL_SUCCESS || vec.empty()) return -1;
*type = static_cast<char*>(ValkeyModule_Alloc(vec[0].length() + 1));
*len = vec[0].length();
snprintf(*type, *len + 1, "%s", vec[0].c_str());
return 0;
}
int get_json_value(ValkeyModuleCtx *ctx, const char *keyname, const size_t key_len, const char *path,
char **value, size_t *len) {
*value = nullptr;
*len = 0;
JDocument *doc = get_json_document(ctx, keyname, key_len);
if (doc == nullptr) return -1;
rapidjson::StringBuffer buf;
JsonUtilCode rc = dom_get_value_as_str(doc, path, nullptr, buf, false);
if (rc != JSONUTIL_SUCCESS) return -1;
*len = buf.GetLength();
*value = static_cast<char*>(ValkeyModule_Alloc(*len + 1));
snprintf(*value, *len + 1, "%s", buf.GetString());
return 0;
}
int get_json_values_and_types(ValkeyModuleCtx *ctx, const char *keyname, const size_t key_len, const char **paths,
const int num_paths, char ***values, size_t **lengths, char ***types, size_t **type_lengths) {
ValkeyModule_Assert(values != nullptr);
ValkeyModule_Assert(lengths != nullptr);
*values = nullptr;
*lengths = nullptr;
if (types != nullptr) *types = nullptr;
if (type_lengths != nullptr) *type_lengths = nullptr;
JDocument *doc = get_json_document(ctx, keyname, key_len);
if (doc == nullptr) return -1;
*values = static_cast<char **>(ValkeyModule_Alloc(num_paths * sizeof(char *)));
*lengths = static_cast<size_t *>(ValkeyModule_Alloc(num_paths * sizeof(size_t)));
memset(*values, 0, num_paths * sizeof(char *));
memset(*lengths, 0, num_paths * sizeof(size_t));
for (int i = 0; i < num_paths; i++) {
rapidjson::StringBuffer buf;
JsonUtilCode rc = dom_get_value_as_str(doc, paths[i], nullptr, buf, false);
if (rc == JSONUTIL_SUCCESS) {
(*lengths)[i] = buf.GetLength();
(*values)[i] = static_cast<char*>(ValkeyModule_Alloc((*lengths)[i] + 1));
snprintf((*values)[i], (*lengths)[i] + 1, "%s", buf.GetString());
} else {
(*values)[i] = nullptr;
}
}
if (types != nullptr) {
ValkeyModule_Assert(type_lengths != nullptr);
*types = static_cast<char **>(ValkeyModule_Alloc(num_paths * sizeof(char *)));
*type_lengths = static_cast<size_t *>(ValkeyModule_Alloc(num_paths * sizeof(size_t)));
memset(*types, 0, num_paths * sizeof(char *));
memset(*type_lengths, 0, num_paths * sizeof(size_t));
for (int i = 0; i< num_paths; i++) {
jsn::vector<jsn::string> vec;
bool is_v2_path;
JsonUtilCode rc = dom_value_type(doc, paths[i], vec, is_v2_path);
if (rc == JSONUTIL_SUCCESS && !vec.empty()) {
(*type_lengths)[i] = vec[0].length();
(*types)[i] = static_cast<char*>(ValkeyModule_Alloc((*type_lengths)[i] + 1));
snprintf((*types)[i], (*type_lengths)[i] + 1, "%s", vec[0].c_str());
} else {
(*types)[i] = nullptr;
}
}
}
return 0;
}

70
src/json/json_api.h Normal file
View File

@ -0,0 +1,70 @@
/**
* JSON C API
*/
#ifndef VALKEYJSONMODULE_JSON_API_H_
#define VALKEYJSONMODULE_JSON_API_H_
#include <stdlib.h>
typedef struct ValkeyModuleCtx ValkeyModuleCtx;
typedef struct ValkeyModuleKey ValkeyModuleKey;
typedef struct ValkeyModuleString ValkeyModuleString;
#ifdef __cplusplus
extern "C" {
#endif
/**
* Is it a JSON key?
*/
int is_json_key(ValkeyModuleCtx *ctx, ValkeyModuleKey *key);
/**
* Another version of is_json_key, given key name as ValkeyModuleString
*/
int is_json_key2(ValkeyModuleCtx *ctx, ValkeyModuleString *keystr);
/**
* Get the type of the JSON value at the path. The path is expected to point to a single value.
* If multiple values match the path, only the type of the first one is returned.
*
* @type output param, JSON type. The caller is responsible for freeing the memory.
* @len output param, length of JSON type.
* @return 0 - success, 1 - error
*/
int get_json_value_type(ValkeyModuleCtx *ctx, const char *keyname, const size_t key_len, const char *path,
char **type, size_t *len);
/**
* Get serialized JSON value at the path. The path is expected to point to a single value.
* If multiple values match the path, only the first one is returned.
*
* @value output param, serialized JSON string. The caller is responsible for freeing the memory.
* @len output param, length of JSON string.
* @return 0 - success, 1 - error
*/
int get_json_value(ValkeyModuleCtx *ctx, const char *keyname, const size_t key_len, const char *path,
char **value, size_t *len);
/**
* Get serialized JSON values and JSON types at multiple paths. Each path is expected to point to a single value.
* If multiple values match the path, only the first one is returned.
*
* @values Output param, array of JSON strings.
* The caller is responsible for freeing the memory: the array '*values' as well as all the strings '(*values)[i]'.
* @lengths Output param, array of lengths of each JSON string.
* The caller is responsible for freeing the memory: the array '*lengths'.
* @types Output param, array of types as strings. The caller is responsible for freeing the memory.
* The caller is responsible for freeing the memory: the array '*types' as well as all the strings '(*types)[i]'.
* @type_lengths Output param, array of lengths of each type string.
* The caller is responsible for freeing the memory: the array '*type_lengths'.
* @return 0 - success, 1 - error
*/
int get_json_values_and_types(ValkeyModuleCtx *ctx, const char *keyname, const size_t key_len, const char **paths,
const int num_paths, char ***values, size_t **lengths, char ***types, size_t **type_lengths);
#ifdef __cplusplus
}
#endif
#endif

659
src/json/keytable.cc Normal file
View File

@ -0,0 +1,659 @@
#include <stddef.h>
#include <stdint.h>
#include <memory>
#include <sstream>
#include <mutex>
#include <unordered_map>
#include <cstring>
#include <iostream>
extern "C" {
#include <./include/valkeymodule.h>
}
#define KEYTABLE_ASSERT ValkeyModule_Assert
#include <json/keytable.h>
/***************************************************************************************************
*
* The shard implements a hashtable of entries. Each entry consists of a pointer to a unique string.
*
* We implement open addressing using linear probing (see https://en.wikipedia.org/wiki/Linear_probing)
* See the text for details of the search, insert and deletion algorithms.
*
* A hash table is a vector of pointers to KeyTable_Layout objects. Thus for insertion, searching
* and deletion only a single hashing of the key passed in as a parameter is required, no subsequent
* hash computations are done.
*
* The rehash algorithm needs the hash value for each key in order to insert it into the hash table.
* For simplicity, rehashing is done synchronously.
*
* If the new table size is less than 2^19, we store the low order 19 bits of the original hash value in
* the hash-table entry itself (since it's a pointer) and just use that. Thus the only memory associated
* with the new and old hash tables needs to be accessed, reducing the cache footprint.
*
* If the new table size is greater than 2^19, we're out of bits in the pointer so we have to use the
* full value of the hash. That hash value is not recomputed, because we saved it when we generated
* the entry originally (see KeyTable_Layout), so it's just fetched. But fetching of the hash value will
* cause an extra cache miss, further increasing the cost of a rehash for large tables.
*
* There's a trade-off between the number of shards and the size of each shard hashtable. We really want
* to keep the shard hashtable below 2^19 so that rehashes are fast. Thus when a table size grows to
* be larger than 2^19, we put out a warning into the log that performance would suffer and we should
* increase the number of shards. Someday we could hook this up to an alarm.
*
*/
// non-constant so unit tests can control it
size_t MAX_FAST_TABLE_SIZE = PtrWithMetaData<KeyTable_Layout>::METADATA_MASK + 1;
struct KeyTable_Shard {
typedef PtrWithMetaData<KeyTable_Layout> EntryType;
size_t capacity; // Number of entries in table
size_t size; // Number of current entries
size_t bytes; // number of bytes of all current entries
size_t handles; // number of handles outstanding
size_t maxSearch; // Max length of a search, since last read
EntryType *entries; // Array of String Entries
std::mutex mutex; // lock for this shard, mutable for "validate"
uint32_t rehashes; // number of rehashes, since last read
static constexpr size_t MIN_TABLE_SIZE = 4;
//
// This logic implements the optimization that for Fast tables, we just get the low 19-bits of
// the original hash value. Thereby avoiding an extra cache hit to fetch it from the Key itself
//
size_t getHashValueFromEntry(const EntryType& e) const {
size_t hashValue;
if (capacity < MAX_FAST_TABLE_SIZE) {
hashValue = e.getMetaData();
} else {
hashValue = e->getOriginalHash();
KEYTABLE_ASSERT((hashValue & EntryType::METADATA_MASK) == e.getMetaData());
}
return hashValue;
}
void makeTable(const KeyTable& t, size_t newCapacity) {
newCapacity = std::max(newCapacity, MIN_TABLE_SIZE);
KEYTABLE_ASSERT(newCapacity != capacity); // oops full or empty.
capacity = newCapacity;
entries = new (t.malloc(capacity * sizeof(EntryType))) EntryType[capacity];
}
KeyTable_Shard() : mutex() {
capacity = 0;
size = 0;
bytes = 0;
handles = 0;
entries = nullptr;
rehashes = 0;
maxSearch = 0;
}
//
// The only real use for this is in the unit tests.
// We scan the table to make sure all keys are gone.
//
void destroy(KeyTable& t) {
MEMORY_VALIDATE(entries);
for (size_t i = 0; i < capacity; ++i) {
if (entries[i]) {
KEYTABLE_ASSERT(entries[i]->IsStuck());
t.free(&*entries[i]);
}
entries[i].clear();
}
t.free(entries);
entries = nullptr;
}
~KeyTable_Shard() {
KEYTABLE_ASSERT(entries == nullptr);
}
float loadFactor() { return float(size) / float(capacity); }
size_t hashIndex(size_t hash) const { return hash % capacity; }
KeyTable_Layout *insert(KeyTable& t, size_t hsh, const char *ptr, size_t len, bool noescape) {
std::scoped_lock lck(mutex);
while (loadFactor() > t.getFactors().maxLoad) {
//
// Oops, table too full, resize it larger.
//
size_t newSize = capacity + std::max(size_t(capacity * t.getFactors().grow), size_t(1));
resizeTable(t, newSize);
if (newSize >= MAX_FAST_TABLE_SIZE) {
ValkeyModule_Log(nullptr, "warning",
"Fast KeyTable Shard size exceeded, increase "
"json.key-table-num-shards to improve performance");
}
}
size_t ix = hashIndex(hsh);
size_t metadata = hsh & EntryType::METADATA_MASK;
MEMORY_VALIDATE(entries);
for (size_t searches = 0; searches < capacity; ++searches) {
EntryType &entry = entries[ix];
if (!entry) {
//
// Empty, insert it here.
//
handles++;
size++;
bytes += len;
maxSearch = std::max(searches, maxSearch);
KeyTable_Layout *p = KeyTable_Layout::makeLayout(t.malloc, ptr, len, hsh, noescape);
entry = EntryType(p, metadata);
return p;
} else if (entry.getMetaData() == metadata && // Early out, don't hit the cache line....
len == entry->getLength() &&
0 == std::memcmp(ptr, entry->getText(), len)) {
//
// easy case. String already present, just bump the refcount and we're done.
// Use saturating arithmetic so it never fails. If you manage to legitimately
// have a reuse count > 2^29 then you'll never be able to recover the memory
// from that string. But who cares?
//
maxSearch = std::max(searches, maxSearch);
handles++;
if (entry->incrRefCount()) {
t.stuckKeys++;
}
return &*entry;
}
if (++ix >= capacity) {
ix = 0;
}
}
KEYTABLE_ASSERT(false);
return nullptr;
}
//
// forward distance is defined as the number of increments (wrapping around) it takes to go
// from "from" to "to". Accounting for wrap-around
//
size_t forward_distance(size_t from, size_t to) {
size_t result;
if (from <= to) {
result = to - from;
} else {
result = (to + capacity) - from;
}
KEYTABLE_ASSERT(result < capacity);
return result;
}
KeyTable_Layout *clone(KeyTable& t, const KeyTable_Handle& h) {
std::scoped_lock lck(mutex);
handles++;
if (h->incrRefCount()) {
t.stuckKeys++;
}
return const_cast<KeyTable_Layout *>(&*h);
}
void destroyHandle(const KeyTable& t, KeyTable_Handle& h, size_t hsh) {
std::scoped_lock lck(mutex);
handles--;
if (h->decrRefCount() > 0) {
h.clear(); // Kill the handle
return; // Easy case, still referenced.
}
//
// Ok, we need to remove this string from the hashtable.
//
size_t ix = hashIndex(hsh);
MEMORY_VALIDATE(entries);
for (size_t searches = 0; searches < capacity; ++searches) {
if (&*entries[ix] == &*h) {
//
// Found it!!!
// Update stats, nuke the handle and recover the space.
//
KEYTABLE_ASSERT(entries[ix].getMetaData() == (hsh & EntryType::METADATA_MASK));
KEYTABLE_ASSERT(entries[ix]->getRefCount() == 0);
KEYTABLE_ASSERT(size > 0);
KEYTABLE_ASSERT(bytes >= h->getLength());
bytes -= h->getLength();
size--;
h.theHandle->poisonOriginalHash();
t.free(&*h.theHandle);
h.clear(); // Kill the handle
entries[ix].clear();
//
// Now reestablish the invariant of the algorithm by scanning forward until
// we hit another empty cell. While we're scanning we may have to move keys down
// into the newly freed slot.
//
size_t empty_ix = ix; // Remember where the empty slot is.
if (++ix >= capacity) ix = 0; // Next entry
while (entries[ix]) {
KEYTABLE_ASSERT(!entries[empty_ix]);
KEYTABLE_ASSERT(empty_ix != ix);
searches++;
//
// This non-empty key might have to be moved down to the empty slot.
// That happens if the forward_distance of the empty slot is less than
// The forward_distance of the current slot to the native slot for this key.
//
size_t nativeSlot = hashIndex(getHashValueFromEntry(entries[ix]));
if (forward_distance(nativeSlot, ix) > forward_distance(nativeSlot, empty_ix)) {
// Yes, this key can be moved.
entries[empty_ix] = entries[ix];
entries[ix].clear();
empty_ix = ix;
}
if (++ix >= capacity) ix = 0;
}
maxSearch = std::max(searches, maxSearch);
//
// Having removed an entry, check for rehashing
//
if (loadFactor() < t.getFactors().minLoad && capacity > MIN_TABLE_SIZE) {
size_t reduction = std::max(size_t(capacity * t.getFactors().shrink), size_t(1));
resizeTable(t, capacity - reduction);
}
return;
}
if (++ix >= capacity) {
ix = 0;
}
}
KEYTABLE_ASSERT(false); // Not found ????
}
void resizeTable(const KeyTable& t, size_t newSize) {
uint64_t startTime = ValkeyModule_Milliseconds();
if (capacity == newSize) return; // Nothing to do.
KEYTABLE_ASSERT(newSize >= size); // Otherwise it won't fit.
rehashes++;
MEMORY_VALIDATE(entries);
EntryType *oldEntries = entries;
size_t oldCapacity = capacity;
makeTable(t, newSize);
for (size_t i = 0; i < oldCapacity; ++i) {
if (oldEntries[i]) {
//
// Found valid entry, Compute hash to see where it goes.
//
EntryType& oldEntry = oldEntries[i];
KEYTABLE_ASSERT(oldEntry->getRefCount() > 0);
size_t ix = hashIndex(getHashValueFromEntry(oldEntry));
for (size_t searches = 0; searches < capacity; ++searches) {
if (!entries[ix]) {
//
// Empty, insert it.
//
entries[ix] = oldEntry;
maxSearch = std::max(searches, maxSearch);
goto nextOldEntry;
}
if (++ix >= capacity) ix = 0;
}
KEYTABLE_ASSERT(false); // can't fail if
}
nextOldEntry:{}
}
t.free(oldEntries);
uint64_t duration = ValkeyModule_Milliseconds() - startTime;
if (duration == 0) duration = 1;
uint64_t keys_per_second = (size / duration) * 1000;
ValkeyModule_Log(nullptr, "notice",
"Keytable Resize to %zu completed in %lu ms (%lu / sec)",
capacity, duration, keys_per_second);
}
//
// Validate all of the entries and the counters. Unit test stuff.
//
std::string validate(const KeyTable& t, size_t shardNumber) const {
std::scoped_lock lck(const_cast<std::mutex&>(this->mutex)); // Cheat on the mutex
size_t this_refs = 0;
size_t this_size = 0;
size_t this_bytes = 0;
for (size_t i = 0; i < capacity; ++i) {
EntryType e = entries[i];
if (e) {
this_size++;
this_refs += e->getRefCount();
this_bytes += e->getLength();
size_t orig_hash = t.hash(e->getText(), e->getLength());
size_t correct_metadata = orig_hash & EntryType::METADATA_MASK;
size_t nativeIx = hashIndex(orig_hash);
// Validate the metadata field
if (e.getMetaData() != correct_metadata) {
std::ostringstream os;
os << "Found bad metadata in slot " << i << " Metadata:" << e.getMetaData()
<< " Where it should be: " << correct_metadata << " Hash:" << orig_hash
<< " TableSize:" << capacity;
return os.str();
}
//
// Check the Invariant. If this entry isn't in its original slot (hashIndex) then
// none of the locations between the original slot and this one may be empty.
//
for (size_t ix = nativeIx; ix != i;) {
if (!entries[ix]) {
// Error
std::ostringstream os;
os << "Found invalid empty location at slot " << ix << " While validating"
<< " key in slot " << i << " From NativeSlot:" << nativeIx
<< " TableSize:" << capacity;
return os.str();
}
if (++ix >= capacity) ix = 0;
}
}
}
// compare the counts. The summed refcounts only match handle counts if no stuck strings
if (this_size != size ||
(t.stuckKeys == 0 ? this_refs != handles : false) ||
this_bytes != bytes) {
std::ostringstream os;
os << "Count mismatch for shard: " << shardNumber << " Capacity:" << capacity
<< " Handles:" << handles << " sum(refcounts):" << this_refs
<< " Size:" << size << " this_size:" << this_size
<< " Bytes:" << bytes << " this_bytes:" << this_bytes;
return os.str();
}
return std::string(); // Empty means no failure.
}
std::string validate_counts(std::unordered_map<const KeyTable_Layout *, size_t>& counts) const {
std::string result;
std::scoped_lock lck(const_cast<std::mutex&>(this->mutex)); // Cheat on the mutex
for (size_t i = 0; i < capacity; ++i) {
EntryType e = entries[i];
if (e) {
if (counts[&*e] != e->getRefCount()) {
std::ostringstream os;
os
<< "Found bad count for key: " << e->getText()
<< " Found: " << e->getRefCount()
<< " Expected:" << counts[&*e]
<< "\n";
result += os.str();
} else {
counts.erase(&*e);
}
}
}
return result; // Empty means no failures found.
}
// Add our stats to the total so far.
void updateStats(KeyTable::Stats& s) {
std::scoped_lock lck(mutex);
s.size += size;
s.bytes += bytes;
s.handles += handles;
s.maxTableSize = std::max(s.maxTableSize, capacity);
s.minTableSize = std::min(s.minTableSize, capacity);
s.totalTable += capacity;
s.rehashes += rehashes;
s.maxSearch = std::max(s.maxSearch, maxSearch);
//
// Reset the counters
//
maxSearch = 0;
rehashes = 0;
}
// Add our stats
void updateLongStats(KeyTable::LongStats& s, size_t topN) {
std::scoped_lock lck(mutex);
size_t thisRun = 0;
for (size_t i = 0; i < capacity; ++i) {
if (entries[i]) {
thisRun++;
} else if (thisRun != 0) {
s.runs[thisRun]++;
while (s.runs.size() > topN) {
s.runs.erase(s.runs.begin());
}
thisRun = 0;
}
}
}
};
/*
* Setup the KeyTable itself.
*/
KeyTable::KeyTable(const Config& cfg) :
malloc(cfg.malloc),
free(cfg.free),
hash(cfg.hash),
numShards(cfg.numShards),
stuckKeys(0)
{
KEYTABLE_ASSERT(numShards > 0);
KEYTABLE_ASSERT(malloc && free && hash);
shards = new(malloc(numShards * sizeof(KeyTable_Shard))) KeyTable_Shard[numShards];
for (size_t i = 0; i < numShards; ++i) shards[i].makeTable(*this, 1);
KEYTABLE_ASSERT(!isValidFactors(factors));
}
KeyTable::~KeyTable() {
MEMORY_VALIDATE(shards);
for (size_t i = 0; i < numShards; ++i) {
shards[i].destroy(*this);
shards[i].~KeyTable_Shard();
}
free(shards);
shards = nullptr;
}
std::string KeyTable::validate() const {
std::string s;
for (size_t i = 0; i < numShards; ++i) {
s += shards[i].validate(*this, i);
}
return s;
}
std::string KeyTable::validate_counts(std::unordered_map<const KeyTable_Layout *, size_t>& counts) const {
std::string result;
result = validate();
if (!result.empty()) return result;
//
// Now, we need to double compare the current keytable against the counts array.
// We scan the KeyTables and lookup each handle as we go, erasing it from the input counts map.
// If at the end, there are any counts entries left, then we definitely have a problem.
//
for (size_t i = 0; i < numShards; ++i) {
result += shards[i].validate_counts(counts);
}
if (!result.empty()) return result;
//
// Now validate we found everything
//
if (!counts.empty()) {
for (auto& c : counts) {
std::ostringstream os;
os << "Lingering Handle found: " << c.first->getText() << " Count:" << c.second << "\n";
result += os.str();
}
}
return result;
}
KeyTable::Stats KeyTable::getStats() const {
Stats s{};
//
// Global stats
//
s.stuckKeys = stuckKeys;
s.factors = getFactors();
//
// Now sum up the per-shard stats
//
for (size_t i = 0; i < numShards; ++i) {
shards[i].updateStats(s);
}
return s;
}
KeyTable::LongStats KeyTable::getLongStats(size_t topN) const {
LongStats s;
//
// Now sum up the per-shard stats
//
for (size_t i = 0; i < numShards; ++i) {
shards[i].updateLongStats(s, topN);
}
return s;
}
/*
* Take 19 bits from hash, avoid the low end of the hash value as this is used for the per-shard index.
*/
size_t KeyTable::shardNumberFromHash(size_t hash) {
return (hash >> 40) % numShards;
}
size_t KeyTable::hashcodeFromHash(size_t hash) {
return hash & KeyTable_Handle::MAX_HASHCODE;
}
/*
* Upsert a string, returns a handle for this insertion.
*
* This function hashes the string and dispatches the operation to the appropriate shard.
*/
KeyTable_Handle KeyTable::makeHandle(const char *ptr, size_t len, bool noescape) {
size_t hsh = hash(ptr, len);
size_t shardNum = shardNumberFromHash(hsh);
KeyTable_Layout *s = shards[shardNum].insert(*this, hsh, ptr, len, noescape);
return KeyTable_Handle(s, hashcodeFromHash(hsh));
}
/*
* Clone an existing handle
*/
KeyTable_Handle KeyTable::clone(const KeyTable_Handle& h) {
size_t hsh = hash(h.GetString(), h.GetStringLength());
size_t shardNum = shardNumberFromHash(hsh);
KeyTable_Layout *s = shards[shardNum].clone(*this, h);
return KeyTable_Handle(s, hashcodeFromHash(hsh));
}
/*
* Destroy a Handle.
*
* While technically, we don't have to hash the string to determine the shard, the shard-level
* destroy operation will need the hash, so we do it here for symmetry.
*/
void KeyTable::destroyHandle(KeyTable_Handle& h) {
if (!h) return; // Empty
size_t hsh = hash(h.GetString(), h.GetStringLength());
KEYTABLE_ASSERT(!h->isPoisoned());
KEYTABLE_ASSERT(hsh == h->getOriginalHash());
size_t shardNum = shardNumberFromHash(hsh);
shards[shardNum].destroyHandle(*this, h, hsh);
}
void KeyTable::setFactors(const Factors& f) {
KEYTABLE_ASSERT(!isValidFactors(f));
// Grab all of the locks to ensure consistency
for (size_t i = 0; i < numShards; ++i) {
shards[i].mutex.lock();
}
factors = f;
for (size_t i = 0; i < numShards; ++i) {
shards[i].mutex.unlock();
}
}
const char *KeyTable::isValidFactors(const Factors& f) {
//
// first the easy ones....
//
if (f.minLoad <= 0) return "minLoad <= 0.0";
if (f.maxLoad > 1.0f) return "maxLoad > 1.0";
if (f.minLoad >= f.maxLoad) return "minLoad >= maxLoad";
if (f.grow <= 0) return "Grow <= 0.0";
if (f.shrink <= 0) return "Shrink <= 0.0";
//
// The shrink factor requires additional validation because we want to make sure that
// rehash down will always succeed, i.e., you can't shrink TOO much or you're toast.
// (because it won't fit ;-))
//
if (f.shrink > (1.0f - f.minLoad)) return "Shrink too large";
return nullptr; // We're good !!!
}
/******************************************************************************************
* Implement KeyTable_Layout
*
* Efficiently store three quantities in sequential memory: RefCount, Length, Text[0..Length-1]
*
* We use either 1, 2, 3 or 4 bytes to store the length.
*/
// Maximum legal refcount. 2^29-1
static uint32_t MAX_REF_COUNT = 0x1FFFFFFF;
bool KeyTable_Layout::IsStuck() const {
return refCount >= MAX_REF_COUNT;
}
bool KeyTable_Layout::incrRefCount() const {
if (IsStuck()) {
return true; // Saturated
} else {
refCount++;
return false;
}
}
size_t KeyTable_Layout::decrRefCount() const {
KEYTABLE_ASSERT(refCount > 0);
if (!IsStuck()) refCount--;
return refCount;
}
size_t KeyTable_Layout::getLength() const {
// Length is stored in little-endian format
size_t len = 0;
for (size_t i = 0; i <= lengthBytes; ++i) {
len |= *reinterpret_cast<const uint8_t *>(bytes+i) << (i * 8);
}
return len;
}
const char *KeyTable_Layout::getText() const {
return bytes + lengthBytes + 1;
}
KeyTable_Layout *KeyTable_Layout::makeLayout(void *(*malloc)(size_t), const char *ptr, size_t len,
size_t hash, bool noescape) {
size_t lengthBytes = (len <= 0xFF) ? 1 :
(len <= 0xFFFF) ? 2:
(len <= 0xFFFFFF) ? 3:
(len <= 0xFFFFFFFF) ? 4 :
0;
KeyTable_Layout *p = reinterpret_cast<KeyTable_Layout*>(malloc(
sizeof(KeyTable_Layout) + lengthBytes + len));
p->original_hash = hash;
p->noescapeFlag = noescape;
p->refCount = 1;
p->lengthBytes = lengthBytes - 1;
// store the length in little-endian format
for (size_t i = 0; i < lengthBytes; ++i) {
p->bytes[i] = len >> (8 * i);
}
std::memcpy(p->bytes + lengthBytes, ptr, len);
return p;
}
// Unit test only.
void KeyTable_Layout::setMaxRefCount(uint32_t maxRefCount) {
KEYTABLE_ASSERT(sizeof(KeyTable_Layout) == 5 + 8);
KEYTABLE_ASSERT(maxRefCount <= MAX_REF_COUNT); // can only shrink it.
MAX_REF_COUNT = maxRefCount;
}

386
src/json/keytable.h Normal file
View File

@ -0,0 +1,386 @@
#ifndef _KEYTABLE_H
#define _KEYTABLE_H
/************************************************************************************************
*
* The key table. This thread-safe object implements unique use-counted immutable strings.
*
* This object is a repository of immutable strings. You put a string into the repository
* and you get back an immutable handle (8-bytes). The handle can be cheaply de-referenced to yield the
* underlying string text (when needed). When you're done with the handle you give it back to the
* string table. So far, nothing special.
*
* The key table maintains a reference count for each string in the table AND it guarantees that
* each string in the table is unique. Further, two insertions of the same string will yield the
* same handle, meaning that once you've converted a string into a handle you can do equality
* comparisons on other strings simply by comparing the handles for equality.
*
* After initialization, there are only two operations on the global hashtable.
*
* (1) Insert string, return handle. (string is copied, the caller can discard his memory)
* (2) discard handle.
*
* Both operations are thread-safe and the handles are NOT locked to a thread.
* The handle represents a resource allocation within the table and thus every call to (1) must
* eventually have a call to (2) to release the resource (i.e., decrement the refcount)
*
* **********************************************************************************************
*/
/*
* IMPLEMENTATION
*
* Each unique string (with it's metadata, length & refcount) is stored in a separately malloc'ed
* chunk of memory. The handle contains a pointer to this data. Thus having obtained a handle,
* access to the underlying string is trivially cheap. Since the handle and the string itself are
* immutable, no thread locking need be done to access the data.
*
* A separate data structure (table & shards) contains a mapping table that implements the raw
* API to convert a string to a handle. That mapping table is sharded with each shard being
* protected by a mutex. A string is received and hashed. The hashcode selects the shard for that
* string. The shard is locked, the mapping table is consulted to locate a previous copy of the
* string. If the string is found, the refcount is incremented and a handle is constructed from
* the existing pointer and shard number. If the string isn't found, a new malloc string is created,
* the mapping is updated and a handle is constructed and returned.
*
* The mapping is implemented as a hashtable using linear hashing. Each hash table entry is simply
* the malloc'ed pointer and 19-bits of hash code. Various conditions can cause a rehashing event,
* rehashing is always done as a single operation on the entire hashtable while the mutex is held,
* i.e., there is no incremental re-hashing. This makes it very easy to ensure multi-thread correctness.
* Worst-case CPU loads due to rehashing are limited because the size of a shard hashtable is itself
* limited to 2^19 entries. You vary the number of shards to handle the worst-case number of strings
* in the table.
*
* The refcount for a string is currently fixed at 30-bits. Increment and decrement of the refcount
* is done with saturating arithmetic, meaning that if a string ever hits the maximum refcount it
* will never be deleted from the table. This isn't considered to be a problem.
*
*/
#include <atomic>
#include <string>
#include <ostream>
#include <string_view>
#include <map>
#include <unordered_map>
#include "json/alloc.h"
#ifndef KEYTABLE_ASSERT
#define KEYTABLE_ASSERT(x) RAPIDJSON_ASSERT(x)
#endif
/*
* This class implements a pointer with up to 19 bits of additional metadata. On x86_64 and Aarch-64
* There are 16 bits at the top of the pointer that are unused (guaranteed to be zero) and we assume
* that the pointer references malloc'ed memory with a minimum of 8-byte alignment, guaranteeing that
* another 3 bits of usable storage.
*
* Other versions of this class could be implemented for systems that don't meet the requirements
* above and simply store the metadata adjacent to a full pointer (i.e., 32-bit systems).
*
*/
template<typename T>
class PtrWithMetaData {
public:
enum { METADATA_MASK = (1 << 19)-1 }; // Largest value that fits.
const T& operator*() const { return *getPointer(); }
const T* operator->() const { return getPointer(); }
T& operator*() { return *getPointer(); }
T* operator->() { return getPointer(); }
//
// It's "C"-ism, that you test pointers for null/!null by doing "if (ptr)" or "if (!ptr)"
// C++ considers that as a conversion to a boolean. Which is implemented by this operator
// To be clear, we include the Metadata in the comparison.
//
operator bool() const { return bits != 0; } // if (PtrWithMetaData) invokes this operator
size_t getMetaData() const { return ror(bits, 48) & METADATA_MASK; }
void setMetaData(size_t metadata) {
KEYTABLE_ASSERT(0 == (metadata & ~METADATA_MASK));
bits = (bits & PTR_MASK) | ror(metadata, 16);
}
PtrWithMetaData() : bits(0) {}
PtrWithMetaData(T *ptr, size_t metadata) {
KEYTABLE_ASSERT(0 == (~PTR_MASK & reinterpret_cast<size_t>(ptr)));
KEYTABLE_ASSERT(0 == (metadata & ~METADATA_MASK));
bits = reinterpret_cast<size_t>(ptr) | ror(metadata, 16);
}
void clear() { bits = 0; }
//
// Comparison operations also include the metadata
//
bool operator==(const PtrWithMetaData& rhs) const { return bits == rhs.bits; }
bool operator!=(const PtrWithMetaData& rhs) const { return bits != rhs.bits; }
friend std::ostream& operator<<(std::ostream& os, const PtrWithMetaData& ptr) {
return os << "Ptr:" << reinterpret_cast<const void *>(&*ptr) << " MetaData:" << ptr.getMetaData();
}
void swap(PtrWithMetaData& rhs) { std::swap<size_t>(bits, rhs.bits); }
private:
size_t bits;
T* getPointer() const { return MEMORY_VALIDATE<T>(reinterpret_cast<T *>(bits & PTR_MASK)); }
// Circular rotate right (count <= 64)
static constexpr size_t ror(size_t v, unsigned count) {
return (v >> count) | (v << (64-count));
}
static const size_t PTR_MASK = ~ror(METADATA_MASK, 16);
};
//
// This is the struct that's accessed by dereferencing a KeyTable_Handle.
// Normal users should only look at len and text fields -- and consider them immutable.
// For Normal users, the only statement about refcount is that it will be non-zero as long as
// any handle exists.
//
// Privileged unit tests look at the refcount field also...
//
struct KeyTable_Layout {
//
// Create a string layout. allocates some memory
//
static KeyTable_Layout *makeLayout(void *(*malloc)(size_t), const char *ptr,
size_t len, size_t hash, bool noescape);
//
// Interrogate existing layout
//
size_t getRefCount() const { return refCount; }
size_t getLength() const;
const char *getText() const;
bool IsStuck() const;
bool getNoescape() const { return noescapeFlag != 0; }
enum { POISON_VALUE = 0xdeadbeeffeedfeadull };
size_t getOriginalHash() const { return original_hash; }
void poisonOriginalHash() { original_hash = POISON_VALUE; }
bool isPoisoned() const { return original_hash == POISON_VALUE; }
// Unit test
static void setMaxRefCount(uint32_t maxRefCount);
protected:
KeyTable_Layout(); // Nobody gets to create one.
friend class KeyTable_Shard; // Only class allowed to manipulate reference count
bool incrRefCount() const; // true => saturated
size_t decrRefCount() const; // returns current count
size_t original_hash; // Remember original hash
mutable uint32_t refCount:29; // Ref count.
uint32_t noescapeFlag:1; // String doesn't need to be escaped
uint32_t lengthBytes:2; // 0, 1, 2 or 3 => 1, 2, 3 or 4 bytes of length
char bytes[1]; // length bytes + text bytes
} __attribute__((packed)); // Don't let compiler round size of up 8 bytes.
struct KeyTable_Handle {
/***************************** Public Handle Interface *******************************/
//
// get a pointer to the text of the string. This pointer has the same lifetime as the
// string_table_handle object itself.
//
const KeyTable_Layout& operator*() const { return *theHandle; }
const KeyTable_Layout* operator->() const { return &*theHandle; }
const char *GetString() const { return theHandle->getText(); }
size_t GetStringLength() const { return theHandle->getLength(); }
const std::string_view GetStringView() const
{ return std::string_view(theHandle->getText(), theHandle->getLength()); }
size_t GetHashcode() const { return theHandle.getMetaData(); }
bool IsNoescape() const { return theHandle->getNoescape(); }
enum { MAX_HASHCODE = PtrWithMetaData<KeyTable_Layout>::METADATA_MASK };
//
// Assignment is only allowed into a empty handle.
//
KeyTable_Handle& operator=(const KeyTable_Handle& rhs) {
KEYTABLE_ASSERT(!theHandle);
theHandle = rhs.theHandle;
const_cast<KeyTable_Handle&>(rhs).theHandle.clear();
return *this;
}
//
// Do assignment into raw storage
//
void RawAssign(const KeyTable_Handle& rhs) {
theHandle = rhs.theHandle;
const_cast<KeyTable_Handle&>(rhs).theHandle.clear();
}
//
// move semantics are allowed
//
KeyTable_Handle(KeyTable_Handle&& rhs) {
theHandle = rhs.theHandle;
rhs.theHandle.clear();
}
//
// Comparison
//
bool operator==(const KeyTable_Handle& rhs) const { return theHandle == rhs.theHandle; }
bool operator!=(const KeyTable_Handle& rhs) const { return theHandle != rhs.theHandle; }
operator bool() const { return bool(theHandle); }
friend std::ostream& operator<<(std::ostream& os, const KeyTable_Handle& h) {
return os << "Handle:" << reinterpret_cast<const void *>(&*(h.theHandle))
<< " Shard:" << h.theHandle.getMetaData()
<< " RefCount: " << h->getRefCount()
<< " : " << h.GetStringView();
}
KeyTable_Handle() : theHandle() {}
~KeyTable_Handle() { KEYTABLE_ASSERT(!theHandle); }
void Swap(KeyTable_Handle& rhs) {
theHandle.swap(rhs.theHandle);
}
private:
friend class KeyTable;
friend struct KeyTable_Shard;
KeyTable_Handle(KeyTable_Layout *ptr, size_t hashCode) : theHandle(ptr, hashCode) {}
void clear() { theHandle.clear(); }
PtrWithMetaData<KeyTable_Layout> theHandle; // The only actual data here.
};
/*
* This is the core hashtable, it's invisible externally
*/
struct KeyTable_Shard;
struct KeyTable {
/*************************** External Table Interface *********************************/
enum { MAX_SHARDS = KeyTable_Handle::MAX_HASHCODE, MIN_SHARDS = 1 };
//
// Stuff to create a table. These get copied and can't be changed without
// recreating the entire table.
//
struct Config {
void *(*malloc)(size_t); // Use this to allocate memory
void (*free)(void*); // Use this to free memory
size_t (*hash)(const char *, size_t); // Hash function for strings
size_t numShards; // Number of shards to create
};
//
// Construct a table.
//
explicit KeyTable(const Config& cfg);
~KeyTable();
//
// Make a handle for this string. The target string is copied when necessary.
//
KeyTable_Handle makeHandle(const char *ptr, size_t len, bool noescape = false);
KeyTable_Handle makeHandle(const std::string& s, bool noescape = false) {
return makeHandle(s.c_str(), s.length(), noescape);
}
KeyTable_Handle makeHandle(const std::string_view& s, bool noescape = false) {
return makeHandle(s.data(), s.length(), noescape);
}
KeyTable_Handle clone(const KeyTable_Handle& rhs);
//
// Destroy a handle
//
void destroyHandle(KeyTable_Handle &h);
//
// Some of the configuration variables can be changed dynamically.
//
struct Factors {
float minLoad; // LoadFactor() < minLoad => rehash down
float maxLoad; // LoadFactor() > maxLoad => rehash up
float shrink; // % to shrink by
float grow; // % to grow by
Factors() :
// Default Factors for the hash table
minLoad(0.25f), // minLoad => .25
maxLoad(0.85f), // maxLoad => targets O(8) searches [see wikipedia]
shrink(0.5f), // shrink, remove 1/2 of elements.
grow(1.0f) // Grow by 100%
{}
};
//
// Get the current configuration
//
const Factors& getFactors() const { return factors; }
//
// Query if this set of factors is valid.
// returns: NULL, If the factors are valid. Otherwise an error string
// This is used to validate a set of factors before setting them.
//
static const char *isValidFactors(const Factors& f);
//
// Change to these factors if valid. This is modestly expensive as it grabs all shard locks
// This will assert if the factors are invalid.
//
void setFactors(const Factors& proposed);
/*
* Stats you can get at any time.
*
* Reading these is O(numShards) which can be expensive
*
* These stats are computed by summing up across the shards. Each shard is locked and
* then it's contribution is added to the running totals. Because of the time skew for
* the reading, there maybe slight inaccuracies in the presence of multi-thread operations.
*/
struct Stats {
size_t size; // Total number of unique strings in table
size_t bytes; // Total bytes of strings
size_t handles; // Number of outstanding handles
size_t maxTableSize; // Largest Shard table
size_t minTableSize; // Smallest Shard table
size_t totalTable; // sum of table sizes
size_t stuckKeys; // Number of strings that have hit the refcount max.
//
// These counters are reset after being read.
//
size_t maxSearch; // longest search sequence encountered
size_t rehashes; // Number of rehashes
//
// Include a copy of current settable factors. Makes testing easier
//
Factors factors;
};
Stats getStats() const;
//
// Long stats are stats that are VERY expensive to compute and are generally only
// used for debug or unit tests. You can see these in the JSON.DEBUG command provided
// you are coming in from an Admin connection.
//
struct LongStats {
std::map<size_t, size_t> runs; // size of of runs, count of #runs
};
//
// TopN parameter limits size of result to largest N runs.
// Setting N to a relatively small number will reduce the cost of generating the stats.
//
LongStats getLongStats(size_t topN) const;
KeyTable(const KeyTable& rhs) = delete; // no copies
void operator=(const KeyTable& rhs) = delete; // no assignment
std::string validate() const; // Unit testing only
std::string validate_counts(std::unordered_map<const KeyTable_Layout *, size_t>& counts) const; // Debug command
size_t getNumShards() const { return numShards; }
private:
friend class KeyTable_Shard;
size_t shardNumberFromHash(size_t hash);
size_t hashcodeFromHash(size_t hash);
KeyTable_Shard* shards;
void *(*malloc)(size_t); // Use this to allocate memory
void (*free)(void *); // Use this to free memory
size_t (*hash)(const char *, size_t); // Hash function for strings
size_t numShards;
std::atomic<size_t> stuckKeys; // Stuck String count.
Factors factors;
};
extern KeyTable *keyTable; // The singleton
#endif

352
src/json/memory.cc Normal file
View File

@ -0,0 +1,352 @@
#include <stdlib.h>
#include <string.h>
#include <atomic>
#include <iostream>
#include "json/memory.h"
#include "json/dom.h"
extern "C" {
#define VALKEYMODULE_EXPERIMENTAL_API
#include <./include/valkeymodule.h>
}
#define STATIC /* decorator for static functions, remove so that backtrace symbols include these */
void *(*memory_alloc)(size_t size);
void (*memory_free)(void *ptr);
void *(*memory_realloc)(void *orig_ptr, size_t new_size);
size_t (*memory_allocsize)(void *ptr);
bool memoryTrapsEnabled = true;
static std::atomic<size_t> totalMemoryUsage;
size_t memory_usage() {
return totalMemoryUsage;
}
/*
* When Traps are disabled, The following code is used
*/
STATIC void *memory_alloc_without_traps(size_t size) {
void *ptr = ValkeyModule_Alloc(size);
totalMemoryUsage += ValkeyModule_MallocSize(ptr);
return ptr;
}
STATIC void memory_free_without_traps(void *ptr) {
if (!ptr) return;
size_t sz = ValkeyModule_MallocSize(ptr);
ValkeyModule_Assert(sz <= totalMemoryUsage);
totalMemoryUsage -= sz;
ValkeyModule_Free(ptr);
}
STATIC void *memory_realloc_without_traps(void *ptr, size_t new_size) {
if (ptr) {
size_t old_size = ValkeyModule_MallocSize(ptr);
ValkeyModule_Assert(old_size <= totalMemoryUsage);
totalMemoryUsage -= old_size;
}
ptr = ValkeyModule_Realloc(ptr, new_size);
totalMemoryUsage += ValkeyModule_MallocSize(ptr);
return ptr;
}
#define memory_allocsize_without_traps ValkeyModule_MallocSize
//
// Implementation of traps
//
//
// This word of data preceeds the memory allocation as seen by the client.
// The presence of the length is redundant with calling the low-level allocators memory-size function,
// but that function can be fairly expensive, so by duplicating here we optimize the run-time cost.
//
struct trap_prefix {
mutable uint64_t length:40;
mutable uint64_t valid_prefix:24;
enum { VALID = 0xdeadbe, INVALID = 0xf00dad};
static trap_prefix *from_ptr( void *p) { return reinterpret_cast< trap_prefix *>(p) - 1; }
static const trap_prefix *from_ptr(const void *p) { return reinterpret_cast<const trap_prefix *>(p) - 1; }
};
//
// Another word of data is added to end of each allocation. It's set to a known data pattern.
//
struct trap_suffix {
mutable uint64_t valid_suffix;
enum { VALID = 0xdeadfeedbeeff00dull, INVALID = ~VALID };
static trap_suffix *from_prefix(trap_prefix *p) {
return reinterpret_cast<trap_suffix *>(p + 1 + (p->length >> 3));
}
static const trap_suffix *from_prefix(const trap_prefix *p) {
return reinterpret_cast<const trap_suffix *>(p + 1 + (p->length >> 3));
}
};
bool memory_validate_ptr(const void *ptr, bool crashOnError) {
if (!ptr) return true; // Null pointers are valid.
auto prefix = trap_prefix::from_ptr(ptr);
if (prefix->valid_prefix != trap_prefix::VALID) {
if (crashOnError) {
ValkeyModule_Log(nullptr, "error", "Validation Failure memory Corrupted at:%p", ptr);
ValkeyModule_Assert(nullptr == "Validate Prefix Corrupted");
} else {
return false;
}
}
auto suffix = trap_suffix::from_prefix(prefix);
if (suffix->valid_suffix != trap_suffix::VALID) {
if (!crashOnError) return false;
// Dump the first N bytes. Hopefully this might give us a clue what's going wrong....
size_t malloc_size = ValkeyModule_MallocSize(const_cast<void *>(reinterpret_cast<const void *>(prefix)));
ValkeyModule_Assert(malloc_size >= (sizeof(trap_prefix) + sizeof(trap_suffix)));
size_t available_size = malloc_size - (sizeof(trap_prefix) + sizeof(trap_suffix));
size_t dump_size = available_size > 256 ? 256 : available_size;
ValkeyModule_Log(nullptr, "error", "Validation Failure memory overrun @%p size:%zu", ptr, available_size);
auto data = static_cast<const void * const*>(ptr);
while (dump_size > (4 * sizeof(void *))) {
ValkeyModule_Log(nullptr, "error", "Memory[%p]: %p %p %p %p",
static_cast<const void *>(data), data[0], data[1], data[2], data[3]);
data += 4;
dump_size -= 4 * sizeof(void *);
}
while (dump_size) {
ValkeyModule_Log(nullptr, "error", "Memory[%p]: %p",
static_cast<const void *>(data), data[0]);
data++;
dump_size -= sizeof(void *);
}
ValkeyModule_Assert(nullptr == "Validate Suffix Corrupted");
}
return true;
}
STATIC void *memory_alloc_with_traps(size_t size) {
size_t requested_bytes = ~7 & (size + 7); // Round up
size_t alloc_bytes = requested_bytes + sizeof(trap_prefix) + sizeof(trap_suffix);
auto prefix = reinterpret_cast<trap_prefix *>(ValkeyModule_Alloc(alloc_bytes));
totalMemoryUsage += ValkeyModule_MallocSize(prefix);
prefix->valid_prefix = trap_prefix::VALID;
prefix->length = requested_bytes;
auto suffix = trap_suffix::from_prefix(prefix);
suffix->valid_suffix = trap_suffix::VALID;
return reinterpret_cast<void *>(prefix + 1);
}
STATIC void memory_free_with_traps(void *ptr) {
if (!ptr) return;
memory_validate_ptr(ptr);
auto prefix = trap_prefix::from_ptr(ptr);
prefix->valid_prefix = 0;
size_t sz = ValkeyModule_MallocSize(prefix);
ValkeyModule_Assert(sz <= totalMemoryUsage);
totalMemoryUsage -= sz;
ValkeyModule_Free(prefix);
}
STATIC size_t memory_allocsize_with_traps(void *ptr) {
if (!ptr) return 0;
memory_validate_ptr(ptr);
auto prefix = trap_prefix::from_ptr(ptr);
return prefix->length;
}
//
// Do a realloc, but this is rare, so we do it suboptimally, i.e., with a copy
//
STATIC void *memory_realloc_with_traps(void *orig_ptr, size_t new_size) {
if (!orig_ptr) return memory_alloc_with_traps(new_size);
memory_validate_ptr(orig_ptr);
auto new_ptr = memory_alloc_with_traps(new_size);
memcpy(new_ptr, orig_ptr, memory_allocsize_with_traps(orig_ptr));
memory_free_with_traps(orig_ptr);
return new_ptr;
}
//
// Enable/Disable traps
//
bool memory_traps_control(bool enable) {
if (totalMemoryUsage != 0) {
ValkeyModule_Log(nullptr, "warning",
"Attempt to enable/disable memory traps ignored, %zu outstanding memory.", totalMemoryUsage.load());
return false;
}
if (enable) {
memory_alloc = memory_alloc_with_traps;
memory_free = memory_free_with_traps;
memory_realloc = memory_realloc_with_traps;
memory_allocsize = memory_allocsize_with_traps;
} else {
memory_alloc = memory_alloc_without_traps;
memory_free = memory_free_without_traps;
memory_realloc = memory_realloc_without_traps;
memory_allocsize = memory_allocsize_without_traps;
}
memoryTrapsEnabled = enable;
return true;
}
void memory_corrupt_memory(const void *ptr, memTrapsCorruption_t corruption) {
memory_validate_ptr(ptr);
auto prefix = trap_prefix::from_ptr(ptr);
auto suffix = trap_suffix::from_prefix(prefix);
switch (corruption) {
case CORRUPT_PREFIX:
prefix->valid_prefix = trap_prefix::INVALID;
break;
case CORRUPT_LENGTH:
prefix->length--;
break;
case CORRUPT_SUFFIX:
suffix->valid_suffix = trap_suffix::INVALID;
break;
default:
ValkeyModule_Assert(0);
break;
}
}
void memory_uncorrupt_memory(const void *ptr, memTrapsCorruption_t corruption) {
auto prefix = trap_prefix::from_ptr(ptr);
auto suffix = trap_suffix::from_prefix(prefix);
switch (corruption) {
case CORRUPT_PREFIX:
ValkeyModule_Assert(prefix->valid_prefix == trap_prefix::INVALID);
prefix->valid_prefix = trap_prefix::VALID;
break;
case CORRUPT_LENGTH:
prefix->length++;
break;
case CORRUPT_SUFFIX:
ValkeyModule_Assert(suffix->valid_suffix == trap_suffix::INVALID);
suffix->valid_suffix = trap_suffix::VALID;
break;
default:
ValkeyModule_Assert(0);
break;
}
memory_validate_ptr(ptr);
}
//
// Helper functions for JSON validation
//
// true => Valid.
// false => NOT VALID
//
bool ValidateJValue(JValue &v) {
auto p = v.trap_GetMallocPointer(false);
if (p && !memory_validate_ptr(p, false)) return false;
if (v.IsObject()) {
for (auto m = v.MemberBegin(); m != v.MemberEnd(); ++m) {
if (!ValidateJValue(m->value)) return false;
}
} else if (v.IsArray()) {
for (size_t i = 0; i < v.Size(); ++i) {
if (!ValidateJValue(v[i])) return false;
}
}
return true;
}
//
// Dump a JValue with Redaction and memory Validation.
//
// Typical use case:
//
// std::ostringstream os;
// DumpRedactedJValue(os, <jvalue>);
//
void DumpRedactedJValue(std::ostream& os, const JValue &v, size_t level, int index) {
for (size_t i = 0; i < (3 * level); ++i) os << ' '; // Indent
os << "@" << reinterpret_cast<const void *>(&v) << " ";
if (index != -1) os << '[' << index << ']' << ' ';
if (v.IsDouble()) {
os << "double string of length " << v.GetDoubleStringLength();
if (!IS_VALID_MEMORY(v.trap_GetMallocPointer(false))) {
os << " <*INVALID*>\n";
} else if (v.trap_GetMallocPointer(false)) {
os << " @" << v.trap_GetMallocPointer(false) << "\n";
} else {
os << "\n";
}
} else if (v.IsString()) {
os << "String of length " << v.GetStringLength();
if (!IS_VALID_MEMORY(v.trap_GetMallocPointer(false))) {
os << " <*INVALID*>\n";
} else if (v.trap_GetMallocPointer(false)) {
os << " @" << v.trap_GetMallocPointer(false) << "\n";
} else {
os << "\n";
}
} else if (v.IsObject()) {
os << " Object with " << v.MemberCount() << " Members";
if (!IS_VALID_MEMORY(v.trap_GetMallocPointer(false))) {
os << " *INVALID*\n";
} else {
os << " @" << v.trap_GetMallocPointer(false) << '\n';
index = 0;
for (auto m = v.MemberBegin(); m != v.MemberEnd(); ++m) {
DumpRedactedJValue(os, m->value, level+1, index);
index++;
}
}
} else if (v.IsArray()) {
os << "Array with " << v.Size() << " Members";
if (!IS_VALID_MEMORY(v.trap_GetMallocPointer(false))) {
os << " *INVALID*\n";
} else {
os << " @" << v.trap_GetMallocPointer(false) << "\n";
for (size_t index = 0; index < v.Size(); ++index) {
DumpRedactedJValue(os, v[index], level+1, int(index));
}
}
} else {
os << "<scalar>\n";
}
}
//
// This class creates an ostream to the Valkey Log. Each line of output is a single call to the ValkeyLog function
//
class ValkeyLogStreamBuf : public std::streambuf {
std::string line;
ValkeyModuleCtx *ctx;
const char *level;
public:
ValkeyLogStreamBuf(ValkeyModuleCtx *_ctx, const char *_level) : ctx(_ctx), level(_level) {}
~ValkeyLogStreamBuf() {
if (!line.empty()) {
ValkeyModule_Log(ctx, level, "%s", line.c_str());
}
}
std::streamsize xsputn(const char *p, std::streamsize n) {
for (std::streamsize i = 0; i < n; ++i) {
overflow(p[i]);
}
return n;
}
int overflow(int c) {
if (c == '\n' || c == EOF) {
ValkeyModule_Log(ctx, level, "%s", line.c_str());
line.resize(0);
} else {
line += c;
}
return c;
}
};
void DumpRedactedJValue(const JValue &v, ValkeyModuleCtx *ctx, const char *level) {
ValkeyLogStreamBuf b(ctx, level);
std::ostream buf(&b);
DumpRedactedJValue(buf, v);
}

144
src/json/memory.h Normal file
View File

@ -0,0 +1,144 @@
/**
*/
#ifndef VALKEYJSONMODULE_MEMORY_H_
#define VALKEYJSONMODULE_MEMORY_H_
#include <stddef.h>
#include <vector>
#include <set>
#include <unordered_set>
#include <iostream>
#include <string>
#include <sstream>
//
// Trap implementation
//
// Memory traps are a diagnostic tool intended to catch some categories of memory usage errors.
//
// The Trap system is conceptually a shim layer between the client application and the lower level memory allocator.
// Traps operate by adding to each memory allocation a prefix and a suffix. The prefix and suffix contain known
// data patterns and some internal trap metadata. Subsequent memory operations validate the correctness of the
// prefix and suffix. A special interface is provided to allow any client application to voluntarily request
// memory validation -- presumably before utilizing the underlying memory.
//
// This strategy should catch at least three classes of memory corruption:
//
// (1) double free of memory.
// (2) writes off the end of memory (just the prev/next word, not WAAAAY off the end of memory)
// (3) dangling pointer to previously freed memory (this relies on voluntary memory validation)
//
// Traps can be dynamically enabled/disabled, provided that there is no outstanding memory allocation.
//
//
// All functions in the module (outsize of memory.cc) should use these to allocate memory
// instead of the ValkeyModule_xxxx functions.
//
extern void *(*memory_alloc)(size_t size);
extern void (*memory_free)(void *ptr);
extern void *(*memory_realloc)(void *orig_ptr, size_t new_size);
extern size_t (*memory_allocsize)(void *ptr);
//
// Total memory usage.
//
// (1) Includes dom_alloc memory usage. dom_alloc tracks JSON data that's associated with a document
// (2) Includes KeyTable usage, i.e., JSON data that's shared across documents
// (3) Includes STL library allocations
//
extern size_t memory_usage();
//
// Are traps enabled?
//
inline bool memory_traps_enabled() {
extern bool memoryTrapsEnabled;
return memoryTrapsEnabled;
}
//
// External Interface to traps logic
//
// Enables/Disable traps. This can fail if there's outstanding allocated memory.
//
// return true => operation was successful.
// return false => operation failed (there's outstanding memory)
//
bool memory_traps_control(bool enable);
bool memory_validate_ptr(const void *ptr, bool crashOnError = true);
//
// This version validates memory, but crashes on an invalid pointer
//
template<typename t>
static inline t *MEMORY_VALIDATE(t *ptr, bool validate = true) {
extern bool memoryTrapsEnabled;
if (memoryTrapsEnabled && validate) memory_validate_ptr(ptr, true);
return ptr;
}
//
// This version validates memory, but doesn't crash
//
template<typename t>
static inline bool IS_VALID_MEMORY(t *ptr) {
return memory_validate_ptr(ptr, false);
}
//
// Classes for STL Containers that utilize memory usage and trap logic.
//
namespace jsn
{
//
// Our custom allocator
//
template <typename T> class stl_allocator : public std::allocator<T> {
public:
typedef T value_type;
stl_allocator() = default;
stl_allocator(std::allocator<T>&) {}
stl_allocator(std::allocator<T>&&) {}
template <class U> constexpr stl_allocator(const stl_allocator<U>&) noexcept {}
T *allocate(std::size_t n) { return static_cast<T *>(memory_alloc(n*sizeof(T))); }
void deallocate(T *p, std::size_t n) { (void)n; memory_free(p); }
};
template<class Elm> using vector = std::vector<Elm, stl_allocator<Elm>>;
template<class Key, class Compare = std::less<Key>> using set = std::set<Key, Compare, stl_allocator<Key>>;
template<class Key, class Hash = std::hash<Key>, class KeyEqual = std::equal_to<Key>>
using unordered_set = std::unordered_set<Key, Hash, KeyEqual, stl_allocator<Key>>;
typedef std::basic_string<char, std::char_traits<char>, stl_allocator<char>> string;
typedef std::basic_stringstream<char, std::char_traits<char>, stl_allocator<char>> stringstream;
} // namespace jsn
// custom specialization of std::hash can be injected in namespace std
template<>
struct std::hash<jsn::string>
{
std::size_t operator()(const jsn::string& s) const noexcept {
return std::hash<std::string_view>{}(std::string_view(s.c_str(), s.length()));
}
};
//
// Everything below this line is private to this module, it's here for usage by unit tests
//
typedef enum MEMORY_TRAPS_CORRUPTION {
CORRUPT_PREFIX,
CORRUPT_LENGTH,
CORRUPT_SUFFIX
} memTrapsCorruption_t;
void memory_corrupt_memory(const void *ptr, memTrapsCorruption_t corrupt);
void memory_uncorrupt_memory(const void *ptr, memTrapsCorruption_t corrupt);
#endif // VALKEYJSONMODULE_MEMORY_H_

23
src/json/rapidjson_includes.h Normal file
View File

@ -0,0 +1,23 @@
#ifndef _RAPIDJSON_INCLUDES_H
#define _RAPIDJSON_INCLUDES_H
/*
* This file includes all RapidJSON Files (modified or original). Any RAPIDJSON-global #defines, etc. belong here
*/
#if defined(__amd64__) || defined(__amd64) || defined(__x86_64__) || \
defined(__x86_64) || defined(_M_X64) || defined(_M_AMD64)
#define RAPIDJSON_SSE42 1
#endif
#if defined(__ARM_NEON) || defined(__ARM_NEON__)
#define RAPIDJSON_NEON 1
#endif
#define RAPIDJSON_48BITPOINTER_OPTIMIZATION 1
#include "rapidjson/prettywriter.h"
#include "rapidjson/document.h"
#include <rapidjson/encodings.h>
#endif

2418
src/json/selector.cc Normal file
View File

File diff suppressed because it is too large Load Diff

369
src/json/selector.h Normal file
View File

@ -0,0 +1,369 @@
#ifndef VALKEYJSONMODULE_JSON_SELECTOR_H_
#define VALKEYJSONMODULE_JSON_SELECTOR_H_
#include "json/dom.h"
#include "json/rapidjson_includes.h"
#include <string_view>
struct Token {
enum TokenType {
UNKNOWN = 0,
DOLLAR, DOT, DOTDOT, WILDCARD,
COLON, COMMA, AT, QUESTION_MARK,
LBRACKET, RBRACKET, LPAREN, RPAREN,
SINGLE_QUOTE, DOUBLE_QUOTE,
PLUS, MINUS, DIV, PCT,
EQ, NE, GT, LT, GE, LE, NOT, ASSIGN,
ALPHA, DIGIT, SPACE,
TRUE, FALSE, AND, OR,
SPECIAL_CHAR,
END
};
Token()
: type(Token::UNKNOWN)
, strVal()
{}
TokenType type;
std::string_view strVal;
};
/**
* A helper class that contains a string view and an optional internal string. The caller decides if the view
* is a view of an external string or the internal string.
* If StringViewHelper::str is empty, the underlying string is owned by an external resource.
* Otherwise, the underlying string is owned by StringViewHelper.
*/
struct StringViewHelper {
StringViewHelper() : str(), view() {}
StringViewHelper(const StringViewHelper &svh) {
str = svh.str;
if (str.empty())
view = svh.view;
else
view = std::string_view(str.c_str(), str.length());
}
const std::string_view& getView() const { return view; }
void setInternalString(const jsn::string &s) {
str = s;
view = std::string_view(str.c_str(), str.length());
}
void setInternalView(const std::string_view &v) {
str = jsn::string(v);
view = std::string_view(str.c_str(), str.length());
}
void setExternalView(const std::string_view &sv) {
view = sv;
}
private:
jsn::string str;
std::string_view view;
StringViewHelper& operator=(const StringViewHelper&); // disable assignment operator
};
class Lexer {
public:
Lexer()
: p(nullptr)
, next()
, path(nullptr)
, rdTokens(0)
{}
void init(const char *path);
Token::TokenType peekToken() const;
Token nextToken(const bool skipSpace = false);
const Token& currToken() const { return next; }
bool matchToken(const Token::TokenType type, const bool skipSpace = false);
JsonUtilCode scanInteger(int64_t &val);
JsonUtilCode scanUnquotedMemberName(StringViewHelper &member_name);
JsonUtilCode scanPathValue(StringViewHelper &output);
JsonUtilCode scanDoubleQuotedString(JParser& parser);
JsonUtilCode scanDoubleQuotedString(jsn::stringstream &ss);
JsonUtilCode scanSingleQuotedString(jsn::stringstream &ss);
JsonUtilCode scanSingleQuotedStringAndConvertToDoubleQuotedString(jsn::stringstream &ss);
JsonUtilCode scanNumberInFilterExpr(StringViewHelper &number_sv);
JsonUtilCode scanIdentifier(StringViewHelper &sv);
void skipSpaces();
void unescape(const std::string_view &input, jsn::stringstream &ss);
size_t getRecursiveDescentTokens() { return rdTokens; }
const char *p; // current position in path
Token next;
private:
Lexer(const Lexer &t); // disable copy constructor
Lexer& operator=(const Lexer &rhs); // disable assignment constructor
int64_t scanUnsignedInteger();
const char *path;
size_t rdTokens; // number of recursive descent tokens
};
/**
* A JSONPath parser and evaluator that supports both v2 JSONPath and the legacy path syntax, and operates in either
* READ or WRITE mode. It is named Selector because:
* a) For READ, evaluation means selecting a list of values that match the query.
* b) For WRITE, evaluation means selecting a list of values to be updated and places to insert into.
*
* The selector is designed to work with a vector of values instead of a single value, and support both v1 and v2
* path syntax.
*
* Internally, it maintains two pointers. One points to the current node (value) in the JSON tree. The other points
* to the current position in the path string.
*
* The selector automatically detects if the input path is v1 or v2 syntax, and sets the member isV2Path.
* Member mode indicates READ/INSERT/UPDATE/DELETE mode, which is automatically set based on the entry point method
* being invoked, which is getValues or setValues or deleteValues.
*
* 1. READ mode:
* Selector selector;
* JsonUtilCode rc = selector.getValues(doc, path);
*
* The outcome is a result set (selector.resultSet) that matches the query. Each entry is a (value, valuePath) pair.
*
* 2. WRITE mode:
* 2.1. Insert/Update:
* Selector selector;
* JsonUtilCode rc = selector.setValues(doc, path, new_val);
*
* The outcome is 2 collections:
* a) selector.resultSet: values to update. Each entry is a (value, valuePath) pair.
* b) selector.insertPaths: set of insert paths.
*
* Note that setValues takes care of everything (update/insert). As an option, the caller can inspect these vectors
* for verification purpose.
*
* 2.2. Delete:
* Selector selector;
* JsonUtilCode rc = selector.deleteValues(doc, path, numValsDeleted);
*
* The outcome is selector.resultSet, representing values to delete. Each entry is a (value, valuePath) pair.
*
* NOTE:
* a) Inserting into an array value is not allowed. (That's the job of JSON.ARRINSERT and JSON.ARRAPPEND)
* b) A new key can be appended to an object if and only if the key is the last child in the path.
*/
class Selector {
public:
explicit Selector(bool force_v2_path_behavior = false)
: isV2Path(force_v2_path_behavior)
, root(nullptr)
, node(nullptr)
, nodePath()
, lex()
, maxPathDepth(0)
, currPathDepth(0)
, resultSet()
, insertPaths()
, uniqueResultSet()
, mode(READ)
, isRecursiveSearch(false)
, error(JSONUTIL_SUCCESS)
{}
// ValueInfo - (value, path) pair.
// first: JValue pointer
// second: path to the value, which is in json pointer format.
typedef std::pair<JValue*, jsn::string> ValueInfo;
/**
* Entry point for READ query.
* The outcome is selector.resultSet that matches the query. Each entry is a (value, valuePath) pair.
*/
JsonUtilCode getValues(JValue &root, const char *path);
/**
* Entry point for DELETE.
* The outcome is selector.resultSet that matches the query. Each entry is a (value, valuePath) pair.
*/
JsonUtilCode deleteValues(JValue &root, const char *path, size_t &numValsDeleted);
/**
* Entry point for a single stage INSERT/UPDATE, which commits the operation.
* The outcome is 2 vectors:
* 1) selector.resultSet: values to update. Each entry is a (value, valuePath) pair.
* 2) selector.insertPaths: set of insert paths.
*/
JsonUtilCode setValues(JValue &root, const char *path, JValue &new_val);
/**
* Prepare for a 2-stage INSERT/UPDATE. The 2-stage write splits a write operation into two calls:
* prepareSetValues and commit, where prepareSetValues does not change the Valkey data.
* The purpose of having a 2-stage write is to be able to discard the write operation if
* certain conditions are not satisfied.
*/
JsonUtilCode prepareSetValues(JValue &root, const char *path);
/**
* Commit a 2-stage INSERT/UPDATE.
*/
JsonUtilCode commit(JValue &new_val);
bool isLegacyJsonPathSyntax() const { return !isV2Path; }
bool isSyntaxError(JsonUtilCode code) const;
/**
* Given a list of paths, check if there is at least one path that is v2 JSONPath.
*/
static bool has_at_least_one_v2path(const char **paths, const int num_paths) {
for (int i = 0; i < num_paths; i++) {
if (*paths[i] == '$') return true;
}
return false;
}
bool hasValues() const { return !resultSet.empty(); }
bool hasUpdates() const { return !resultSet.empty(); }
bool hasInserts() const { return !insertPaths.empty(); }
size_t getMaxPathDepth() const { return maxPathDepth; }
const jsn::vector<ValueInfo>& getResultSet() const { return resultSet; }
void getSelectedValues(jsn::vector<JValue*> &values) const;
const jsn::vector<Selector::ValueInfo>& getUniqueResultSet();
void dedupe();
bool isV2Path; // if false, it's legacy syntax
private:
enum Mode {
READ,
INSERT,
UPDATE,
INSERT_OR_UPDATE, // JSON.SET could be insert or update or both
DELETE
};
struct State {
State()
: currNode(nullptr)
, nodePath()
, currPathPtr(nullptr)
, currToken()
, currPathDepth(0)
{}
JValue *currNode;
jsn::string nodePath;
const char *currPathPtr;
Token currToken;
size_t currPathDepth;
};
void snapshotState(State &state) {
state.currNode = node;
state.nodePath = nodePath;
state.currPathPtr = lex.p;
state.currToken = lex.next;
state.currPathDepth = currPathDepth;
}
void restoreState(const State &state) {
node = state.currNode;
nodePath = state.nodePath;
lex.p = state.currPathPtr;
lex.next = state.currToken;
currPathDepth = state.currPathDepth;
}
/***
* Initialize the selector.
*/
JsonUtilCode init(JValue &root, const char *path, const Mode mode);
void resetPointers(JValue &currVal, const char *currPath) {
node = &currVal;
lex.p = currPath;
}
void incrPathDepth() {
currPathDepth++;
maxPathDepth = std::max(maxPathDepth, currPathDepth);
}
void decrPathDepth() {
ValkeyModule_Assert(currPathDepth > 0);
currPathDepth--;
}
/***
* Evaluate the path, which includes parsing and evaluating the path.
*/
JsonUtilCode eval();
JsonUtilCode evalMember(JValue &m, const char *path_start);
JsonUtilCode evalObjectMember(const StringViewHelper &member_name, JValue &val);
JsonUtilCode evalArrayMember(int64_t idx);
JsonUtilCode traverseToObjectMember(const StringViewHelper &member_name);
JsonUtilCode traverseToArrayIndex(int64_t idx);
JsonUtilCode parseSupportedPath();
JsonUtilCode parseRelativePath();
JsonUtilCode parseRecursivePath();
JsonUtilCode parseDotPath();
JsonUtilCode parseBracketPath();
JsonUtilCode parseQualifiedPath();
JsonUtilCode parseQualifiedPathElement();
JsonUtilCode parseKey();
JsonUtilCode parseBracketPathElement();
JsonUtilCode parseWildcardInBrackets();
JsonUtilCode parseNameInBrackets();
JsonUtilCode parseQuotedMemberName(jsn::stringstream &ss);
JsonUtilCode parseUnquotedMemberName(StringViewHelper &name);
JsonUtilCode parseIndexExpr();
JsonUtilCode parseSliceStartsWithColon();
JsonUtilCode parseSliceStartsWithInteger(const int64_t start);
JsonUtilCode parseSliceOrUnionOrIndex();
JsonUtilCode parseEndAndStep(const int64_t start);
JsonUtilCode parseStep(const int64_t start, const int64_t end);
JsonUtilCode parseIndex(int64_t &val);
JsonUtilCode parseFilter();
JsonUtilCode parseFilterExpr(jsn::vector<int64_t> &result);
JsonUtilCode parseTerm(jsn::vector<int64_t> &result);
JsonUtilCode parseFactor(jsn::vector<int64_t> &result);
JsonUtilCode parseMemberName(StringViewHelper &name);
JsonUtilCode parseBracketedMemberName(StringViewHelper &member_name);
JsonUtilCode parseComparisonValue(JValue &v);
JsonUtilCode parseComparisonOp(Token::TokenType &op);
JsonUtilCode swapComparisonOpSide(Token::TokenType &op);
JsonUtilCode processUnionOfMembers(const jsn::vector<jsn::string> &member_names);
JsonUtilCode parseUnionOfIndexes(const int64_t fistIndex);
JsonUtilCode processWildcard();
JsonUtilCode processWildcardKey();
JsonUtilCode processWildcardIndex();
JsonUtilCode parseWildcardFilter();
JsonUtilCode processSubscript(const int64_t idx);
JsonUtilCode processSlice(int64_t start, int64_t end, const int64_t step = 1);
JsonUtilCode processUnion(jsn::vector<int64_t> union_indices);
JsonUtilCode processFilterResult(jsn::vector<int64_t> &result);
JsonUtilCode processArrayContains(const StringViewHelper &member_name, const Token::TokenType op,
const JValue &comparison_value, jsn::vector<int64_t> &result);
JsonUtilCode processComparisonExpr(const bool is_self, const StringViewHelper &member_name,
const Token::TokenType op, const JValue &comparison_value,
jsn::vector<int64_t> &result);
JsonUtilCode processComparisonExprAtIndex(const int64_t idx, const StringViewHelper &member_name,
const Token::TokenType op, const JValue &comparison_value,
jsn::vector<int64_t> &result);
JsonUtilCode processAttributeFilter(const StringViewHelper &member_name, jsn::vector<int64_t> &result);
bool evalOp(const JValue *v, const Token::TokenType op, const JValue &comparison_value);
bool deleteValue(const jsn::string &path);
void vectorUnion(const jsn::vector<int64_t> &v, jsn::vector<int64_t> &r, jsn::unordered_set<int64_t> &set);
void vectorIntersection(const jsn::vector<int64_t> &v1, const jsn::vector<int64_t> &v2,
jsn::vector<int64_t> &result);
JsonUtilCode recursiveSearch(JValue &v, const char *p);
void setError(const JsonUtilCode error_code) { error = error_code; }
JsonUtilCode getError() const { return error; }
JValue *root; // the root value, aka the document
JValue *node; // current node (value) in the JSON tree
jsn::string nodePath; // current node's path, which is in json pointer format
Lexer lex;
size_t maxPathDepth;
size_t currPathDepth;
// resultSet - selected values that match the query. In write mode, these are the source values to update
// or delete.
jsn::vector<ValueInfo> resultSet;
// insertPaths - set of insert paths, which is in json pointer path format.
// Only used for write operations that generate INSERTs.
jsn::unordered_set<jsn::string> insertPaths;
// data structure to assist dedupe
jsn::vector<ValueInfo> uniqueResultSet;
Mode mode;
bool isRecursiveSearch; // if we are doing a recursive search we do not wish to add new fields
JsonUtilCode error; // JSONUTIL_SUCCESS indicates no error
};
#endif

291
src/json/stats.cc Normal file
View File

@ -0,0 +1,291 @@
#include "json/stats.h"
#include <pthread.h>
#include <cstring>
#include <atomic>
#include <string>
#include <iostream>
#include <sstream>
extern "C" {
#define VALKEYMODULE_EXPERIMENTAL_API
#include <./include/valkeymodule.h>
}
#define STATIC /* decorator for static functions, remove so that backtrace symbols include these */
LogicalStats logical_stats; // initialize global variable
// Thread local storage (TLS) key for calculating used memory per thread.
static pthread_key_t thread_local_mem_counter_key;
/* JSON statistics struct.
* Use atomic integers due to possible multi-threading execution of rdb_load and
* also the overhead of atomic operations are negligible.
*/
typedef struct {
std::atomic_ullong used_mem; // global used memory counter
std::atomic_ullong num_doc_keys;
std::atomic_ullong max_depth_ever_seen;
std::atomic_ullong max_size_ever_seen;
std::atomic_ullong defrag_count;
std::atomic_ullong defrag_bytes;
void reset() {
used_mem = 0;
num_doc_keys = 0;
max_depth_ever_seen = 0;
max_size_ever_seen = 0;
defrag_count = 0;
defrag_bytes = 0;
}
} JsonStats;
static JsonStats jsonstats;
// histograms
#define NUM_BUCKETS (11)
static size_t buckets[] = {
0, 256, 1024, 4*1024, 16*1024, 64*1024, 256*1024, 1024*1024,
4*1024*1024, 16*1024*1024, 64*1024*1024, SIZE_MAX
};
// static histogram showing document size distribution
static size_t doc_hist[NUM_BUCKETS];
// dynamic histogram for read operations (JSON.GET and JSON.MGET only)
static size_t read_hist[NUM_BUCKETS];
// dynamic histogram for insert operations (JSON.SET and JSON.ARRINSERT)
static size_t insert_hist[NUM_BUCKETS];
// dynamic histogram for update operations (JSON.SET, JSON.STRAPPEND and JSON.ARRAPPEND)
static size_t update_hist[NUM_BUCKETS];
// dynamic histogram for delete operations (JSON.DEL, JSON.FORGET, JSON.ARRPOP and JSON.ARRTRIM)
static size_t delete_hist[NUM_BUCKETS];
JsonUtilCode jsonstats_init() {
ValkeyModule_Assert(jsonstats.used_mem == 0); // Otherwise you'll lose memory accounting
// Create thread local key. No need to have destructor hook, as the key is created on stack.
if (pthread_key_create(&thread_local_mem_counter_key, nullptr) != 0)
return JSONUTIL_FAILED_TO_CREATE_THREAD_SPECIFIC_DATA_KEY;
jsonstats.reset();
logical_stats.reset();
memset(doc_hist, 0, sizeof(doc_hist));
memset(read_hist, 0, sizeof(read_hist));
memset(insert_hist, 0, sizeof(insert_hist));
memset(update_hist, 0, sizeof(update_hist));
memset(delete_hist, 0, sizeof(delete_hist));
return JSONUTIL_SUCCESS;
}
int64_t jsonstats_begin_track_mem() {
return reinterpret_cast<int64_t>(pthread_getspecific(thread_local_mem_counter_key));
}
int64_t jsonstats_end_track_mem(const int64_t begin_val) {
int64_t end_val = reinterpret_cast<int64_t>(pthread_getspecific(thread_local_mem_counter_key));
return end_val - begin_val;
}
void jsonstats_increment_used_mem(size_t delta) {
// update the atomic global counter
jsonstats.used_mem += delta;
// update the thread local counter
int64_t curr_val = reinterpret_cast<int64_t>(pthread_getspecific(thread_local_mem_counter_key));
pthread_setspecific(thread_local_mem_counter_key, reinterpret_cast<int64_t*>(curr_val + delta));
}
void jsonstats_decrement_used_mem(size_t delta) {
// update the atomic global counter
ValkeyModule_Assert(delta <= jsonstats.used_mem);
jsonstats.used_mem -= delta;
// update the thread local counter
int64_t curr_val = reinterpret_cast<int64_t>(pthread_getspecific(thread_local_mem_counter_key));
pthread_setspecific(thread_local_mem_counter_key, reinterpret_cast<int64_t*>(curr_val - delta));
}
unsigned long long jsonstats_get_used_mem() {
return jsonstats.used_mem;
}
unsigned long long jsonstats_get_num_doc_keys() {
return jsonstats.num_doc_keys;
}
unsigned long long jsonstats_get_max_depth_ever_seen() {
return jsonstats.max_depth_ever_seen;
}
void jsonstats_update_max_depth_ever_seen(const size_t max_depth) {
if (max_depth > jsonstats.max_depth_ever_seen) {
jsonstats.max_depth_ever_seen = max_depth;
}
}
unsigned long long jsonstats_get_max_size_ever_seen() {
return jsonstats.max_size_ever_seen;
}
void jsonstats_update_max_size_ever_seen(const size_t max_size) {
if (max_size > jsonstats.max_size_ever_seen) {
jsonstats.max_size_ever_seen = max_size;
}
}
unsigned long long jsonstats_get_defrag_count() {
return jsonstats.defrag_count;
}
void jsonstats_increment_defrag_count() {
jsonstats.defrag_count++;
}
unsigned long long jsonstats_get_defrag_bytes() {
return jsonstats.defrag_bytes;
}
void jsonstats_increment_defrag_bytes(const size_t amount) {
jsonstats.defrag_bytes += amount;
}
/* Given a size (bytes), find histogram bucket index using binary search.
*/
uint32_t jsonstats_find_bucket(size_t size) {
int lo = 0;
int hi = NUM_BUCKETS; // length of buckets[] is NUM_BUCKETS + 1
while (hi - lo > 1) {
uint32_t mid = (lo + hi) / 2;
if (size < buckets[mid])
hi = mid;
else if (size > buckets[mid])
lo = mid;
else
return mid;
}
return lo;
}
/* Update the static document histogram */
STATIC void update_doc_hist(JDocument *doc, const size_t orig_size, const size_t new_size,
const JsonCommandType cmd_type) {
switch (cmd_type) {
case JSONSTATS_INSERT: {
if (orig_size == 0) {
uint32_t new_bucket = jsonstats_find_bucket(new_size);
doc_hist[new_bucket]++;
dom_set_bucket_id(doc, new_bucket);
} else {
update_doc_hist(doc, orig_size, new_size, JSONSTATS_UPDATE);
}
break;
}
case JSONSTATS_UPDATE: {
if (orig_size != new_size) {
uint32_t orig_bucket = dom_get_bucket_id(doc);
uint32_t new_bucket = jsonstats_find_bucket(new_size);
if (orig_bucket != new_bucket) {
doc_hist[orig_bucket]--;
doc_hist[new_bucket]++;
dom_set_bucket_id(doc, new_bucket);
}
}
break;
}
case JSONSTATS_DELETE: {
uint32_t orig_bucket = dom_get_bucket_id(doc);
if (new_size == 0) {
doc_hist[orig_bucket]--;
} else {
uint32_t new_bucket = jsonstats_find_bucket(new_size);
if (new_bucket != orig_bucket) {
doc_hist[orig_bucket]--;
doc_hist[new_bucket]++;
dom_set_bucket_id(doc, new_bucket);
}
}
break;
}
default:
break;
}
}
void jsonstats_sprint_hist_buckets(char *buf, const size_t buf_size) {
std::ostringstream oss;
oss << "[";
for (size_t i=0; i < NUM_BUCKETS; i++) {
if (i > 0) oss << ",";
oss << buckets[i];
}
oss << ",INF]";
std::string str = oss.str();
ValkeyModule_Assert(str.length() <= buf_size);
memcpy(buf, str.c_str(), str.length());
buf[str.length()] = '\0';
}
STATIC void sprint_hist(size_t *arr, const size_t len, char *buf, const size_t buf_size) {
std::ostringstream oss;
oss << "[";
for (size_t i=0; i < len; i++) {
if (i > 0) oss << ",";
oss << arr[i];
}
oss << "]";
std::string str = oss.str();
ValkeyModule_Assert(str.length() <= buf_size);
memcpy(buf, str.c_str(), str.length());
buf[str.length()] = '\0';
}
void jsonstats_sprint_doc_hist(char *buf, const size_t buf_size) {
sprint_hist(doc_hist, NUM_BUCKETS, buf, buf_size);
}
void jsonstats_sprint_read_hist(char *buf, const size_t buf_size) {
sprint_hist(read_hist, NUM_BUCKETS, buf, buf_size);
}
void jsonstats_sprint_insert_hist(char *buf, const size_t buf_size) {
sprint_hist(insert_hist, NUM_BUCKETS, buf, buf_size);
}
void jsonstats_sprint_update_hist(char *buf, const size_t buf_size) {
sprint_hist(update_hist, NUM_BUCKETS, buf, buf_size);
}
void jsonstats_sprint_delete_hist(char *buf, const size_t buf_size) {
sprint_hist(delete_hist, NUM_BUCKETS, buf, buf_size);
}
void jsonstats_update_stats_on_read(const size_t fetched_val_size) {
uint32_t bucket = jsonstats_find_bucket(fetched_val_size);
read_hist[bucket]++;
}
void jsonstats_update_stats_on_insert(JDocument *doc, const bool is_delete_doc_key, const size_t orig_size,
const size_t new_size, const size_t inserted_val_size) {
if (is_delete_doc_key) jsonstats.num_doc_keys++;
update_doc_hist(doc, orig_size, new_size, JSONSTATS_INSERT);
uint32_t bucket = jsonstats_find_bucket(inserted_val_size);
insert_hist[bucket]++;
jsonstats_update_max_size_ever_seen(new_size);
}
void jsonstats_update_stats_on_update(JDocument *doc, const size_t orig_size, const size_t new_size,
const size_t input_json_size) {
update_doc_hist(doc, orig_size, new_size, JSONSTATS_UPDATE);
uint32_t bucket = jsonstats_find_bucket(input_json_size);
update_hist[bucket]++;
jsonstats_update_max_size_ever_seen(new_size);
}
void jsonstats_update_stats_on_delete(JDocument *doc, const bool is_delete_doc_key, const size_t orig_size,
const size_t new_size, const size_t deleted_val_size) {
update_doc_hist(doc, orig_size, new_size, JSONSTATS_DELETE);
if (is_delete_doc_key) {
ValkeyModule_Assert(jsonstats.num_doc_keys > 0);
jsonstats.num_doc_keys--;
}
uint32_t bucket = jsonstats_find_bucket(deleted_val_size);
delete_hist[bucket]++;
}

136
src/json/stats.h Normal file
View File

@ -0,0 +1,136 @@
/**
* The STATS module's main responsibility is to produce the following metrics:
* 1. Core metrics:
* json_total_memory_bytes: total memory allocated to JSON objects
* json_num_documents: number of document keys in Valkey
* 2. Histograms:
* json_doc_histogram: static histogram showing document size distribution. Value of the i_th element is
* number of documents whose size fall into bucket i.
* json_read_histogram: dynamic histogram for read operations (JSON.GET and JSON.MGET). Value of the i_th
* element is number of read operations with fetched JSON size falling into bucket i.
* json_insert_histogram: dynamic histogram for insert operations (JSON.SET and JSON.ARRINSERT) that either
* insert new documents or insert values into existing documents. Value of the i_th element is number of
* insert operations with inserted values' size falling into bucket i.
* json_update_histogram: dynamic histogram for update operations (JSON.SET, JSON.STRAPPEND and
* JSON.ARRAPPEND). Value of the i_th element is number of update operations with input JSON size falling into
* bucket i.
* json_delete_histogram: dynamic histogram for delete operations (JSON.DEL, JSON.FORGET, JSON.ARRPOP and
* JSON.ARRTRIM). Value of the i_th element is number of delete operations with deleted values' size falling
* into bucket i.
*
* Histogram buckets:
* [0,256), [256,1k), [1k,4k), [4k,16k), [16k,64k), [64k,256k), [256k,1m), [1m,4m), [4m,16m), [16m,64m), [64m,INF).
* Each bucket represents a JSON size range in bytes.
*
* To query metrics, run Valkey command:
* info modules: returns all metrics of the module
* info json_core_metrics: returns core metrics
*/
#ifndef VALKEYJSONMODULE_JSON_STATS_H_
#define VALKEYJSONMODULE_JSON_STATS_H_
#include "json/dom.h"
typedef enum {
JSONSTATS_READ = 0,
JSONSTATS_INSERT,
JSONSTATS_UPDATE,
JSONSTATS_DELETE
} JsonCommandType;
/* Initialize statistics counters and thread local storage (TLS) keys. */
JsonUtilCode jsonstats_init();
/* Begin tracking memory usage.
* @return value of the thread local counter.
*/
int64_t jsonstats_begin_track_mem();
/* End tracking memory usage.
* @param begin_val - previous saved thread local value that is returned from jsonstats_begin_track_memory().
* @return delta of used memory
*/
int64_t jsonstats_end_track_mem(const int64_t begin_val);
/* Get the total memory allocated to JSON objects. */
unsigned long long jsonstats_get_used_mem();
/* The following two methods are invoked by the DOM memory allocator upon every malloc/free/realloc.
* Two memory counters are updated: A global atomic counter and a thread local counter (per thread).
*/
void jsonstats_increment_used_mem(size_t delta);
void jsonstats_decrement_used_mem(size_t delta);
// get counters
unsigned long long jsonstats_get_num_doc_keys();
unsigned long long jsonstats_get_max_depth_ever_seen();
void jsonstats_update_max_depth_ever_seen(const size_t max_depth);
unsigned long long jsonstats_get_max_size_ever_seen();
unsigned long long jsonstats_get_defrag_count();
void jsonstats_increment_defrag_count();
unsigned long long jsonstats_get_defrag_bytes();
void jsonstats_increment_defrag_bytes(const size_t amount);
// updating stats on read/insert/update/delete operation
void jsonstats_update_stats_on_read(const size_t fetched_val_size);
void jsonstats_update_stats_on_insert(JDocument *doc, const bool is_delete_doc_key, const size_t orig_size,
const size_t new_size, const size_t inserted_val_size);
void jsonstats_update_stats_on_update(JDocument *doc, const size_t orig_size, const size_t new_size,
const size_t input_json_size);
void jsonstats_update_stats_on_delete(JDocument *doc, const bool is_delete_doc_key, const size_t orig_size,
const size_t new_size, const size_t deleted_val_size);
// helper methods for printing histograms into C string
void jsonstats_sprint_hist_buckets(char *buf, const size_t buf_size);
void jsonstats_sprint_doc_hist(char *buf, const size_t buf_size);
void jsonstats_sprint_read_hist(char *buf, const size_t buf_size);
void jsonstats_sprint_insert_hist(char *buf, const size_t buf_size);
void jsonstats_sprint_update_hist(char *buf, const size_t buf_size);
void jsonstats_sprint_delete_hist(char *buf, const size_t buf_size);
/* Given a size (bytes), find the histogram bucket index using binary search.
*/
uint32_t jsonstats_find_bucket(size_t size);
/* JSON logical statistics.
* Used for internal tracking of elements for Skyhook Billing.
* Using a similar structure to JsonStats.
* We don't track the logical bytes themselves here as they are tracked by Skyhook Metering.
* We are using size_t to match Valkey Module API for Data Metering.
*/
typedef struct {
std::atomic_size_t boolean_count; // 16 bytes
std::atomic_size_t number_count; // 16 bytes
std::atomic_size_t sum_extra_numeric_chars; // 1 byte per char
std::atomic_size_t string_count; // 16 bytes
std::atomic_size_t sum_string_chars; // 1 byte per char
std::atomic_size_t null_count; // 16 bytes
std::atomic_size_t array_count; // 16 bytes
std::atomic_size_t sum_array_elements; // internal metric
std::atomic_size_t object_count; // 16 bytes
std::atomic_size_t sum_object_members; // internal metric
std::atomic_size_t sum_object_key_chars; // 1 byte per char
void reset() {
boolean_count = 0;
number_count = 0;
sum_extra_numeric_chars = 0;
string_count = 0;
sum_string_chars = 0;
null_count = 0;
array_count = 0;
sum_array_elements = 0;
object_count = 0;
sum_object_members = 0;
sum_object_key_chars = 0;
}
} LogicalStats;
extern LogicalStats logical_stats;
#define DOUBLE_CHARS_CUTOFF 24
#endif // VALKEYJSONMODULE_JSON_STATS_H_

142
src/json/util.cc Normal file
View File

@ -0,0 +1,142 @@
#include "json/util.h"
#include "json/dom.h"
#include "json/alloc.h"
#include <cstring>
#include <string>
#include "json/rapidjson_includes.h"
const char *jsonutil_code_to_message(JsonUtilCode code) {
switch (code) {
case JSONUTIL_SUCCESS:
case JSONUTIL_WRONG_NUM_ARGS:
case JSONUTIL_NX_XX_CONDITION_NOT_SATISFIED:
// only used as code, no message needed
break;
case JSONUTIL_JSON_PARSE_ERROR: return "SYNTAXERR Failed to parse JSON string due to syntax error";
case JSONUTIL_NX_XX_SHOULD_BE_MUTUALLY_EXCLUSIVE:
return "SYNTAXERR Option NX and XX should be mutually exclusive";
case JSONUTIL_INVALID_JSON_PATH: return "SYNTAXERR Invalid JSON path";
case JSONUTIL_INVALID_MEMBER_NAME: return "SYNTAXERR Invalid object member name";
case JSONUTIL_INVALID_NUMBER: return "SYNTAXERR Invalid number";
case JSONUTIL_INVALID_IDENTIFIER: return "SYNTAXERR Invalid identifier";
case JSONUTIL_INVALID_DOT_SEQUENCE: return "SYNTAXERR Invalid dot sequence";
case JSONUTIL_EMPTY_EXPR_TOKEN: return "SYNTAXERR Expression token cannot be empty";
case JSONUTIL_ARRAY_INDEX_NOT_NUMBER: return "SYNTAXERR Array index is not a number";
case JSONUTIL_STEP_CANNOT_NOT_BE_ZERO: return "SYNTAXERR Step in the slice cannot be zero";
case JSONUTIL_INVALID_USE_OF_WILDCARD: return "ERR Invalid use of wildcard";
case JSONUTIL_JSON_PATH_NOT_EXIST: return "NONEXISTENT JSON path does not exist";
case JSONUTIL_PARENT_ELEMENT_NOT_EXIST: return "NONEXISTENT Parent element does not exist";
case JSONUTIL_DOCUMENT_KEY_NOT_FOUND: return "NONEXISTENT Document key does not exist";
case JSONUTIL_NOT_A_DOCUMENT_KEY: return "WRONGTYPE Not a JSON document key";
case JSONUTIL_FAILED_TO_DELETE_VALUE: return "OPFAIL Failed to delete JSON value";
case JSONUTIL_JSON_ELEMENT_NOT_NUMBER: return "WRONGTYPE JSON element is not a number";
case JSONUTIL_JSON_ELEMENT_NOT_BOOL: return "WRONGTYPE JSON element is not a bool";
case JSONUTIL_JSON_ELEMENT_NOT_STRING: return "WRONGTYPE JSON element is not a string";
case JSONUTIL_JSON_ELEMENT_NOT_OBJECT: return "WRONGTYPE JSON element is not an object";
case JSONUTIL_JSON_ELEMENT_NOT_ARRAY: return "WRONGTYPE JSON element is not an array";
case JSONUTIL_VALUE_NOT_NUMBER: return "WRONGTYPE Value is not a number";
case JSONUTIL_VALUE_NOT_STRING: return "WRONGTYPE Value is not a string";
case JSONUTIL_VALUE_NOT_INTEGER: return "WRONGTYPE Value is not an integer";
case JSONUTIL_PATH_SHOULD_BE_AT_THE_END: return "SYNTAXERR Path arguments should be positioned at the end";
case JSONUTIL_COMMAND_SYNTAX_ERROR: return "SYNTAXERR Command syntax error";
case JSONUTIL_MULTIPLICATION_OVERFLOW: return "OVERFLOW Multiplication would overflow";
case JSONUTIL_ADDITION_OVERFLOW: return "OVERFLOW Addition would overflow";
case JSONUTIL_EMPTY_JSON_OBJECT: return "EMPTYVAL Empty JSON object";
case JSONUTIL_EMPTY_JSON_ARRAY: return "EMPTYVAL Empty JSON array";
case JSONUTIL_INDEX_OUT_OF_ARRAY_BOUNDARIES: return "OUTOFBOUNDARIES Array index is out of bounds";
case JSONUTIL_UNKNOWN_SUBCOMMAND: return "SYNTAXERR Unknown subcommand";
case JSONUTIL_FAILED_TO_CREATE_THREAD_SPECIFIC_DATA_KEY:
return "PTHREADERR Failed to create thread-specific data key";
case JSONUTIL_DOCUMENT_SIZE_LIMIT_EXCEEDED:
return "LIMIT Document size limit is exceeded";
case JSONUTIL_DOCUMENT_PATH_LIMIT_EXCEEDED:
return "LIMIT Document path nesting limit is exceeded";
case JSONUTIL_PARSER_RECURSION_DEPTH_LIMIT_EXCEEDED:
return "LIMIT Parser recursion depth is exceeded";
case JSONUTIL_RECURSIVE_DESCENT_TOKEN_LIMIT_EXCEEDED:
return "LIMIT Total number of recursive descent tokens in the query string exceeds the limit";
case JSONUTIL_QUERY_STRING_SIZE_LIMIT_EXCEEDED:
return "LIMIT Query string size limit is exceeded";
case JSONUTIL_CANNOT_INSERT_MEMBER_INTO_NON_OBJECT_VALUE:
return "ERROR Cannot insert a member into a non-object value";
case JSONUTIL_INVALID_RDB_FORMAT:
return "ERROR Invalid value in RDB format";
case JSONUTIL_DOLLAR_CANNOT_APPLY_TO_NON_ROOT: return "SYNTAXERR Dollar sign cannot apply to non-root element";
default: ValkeyModule_Assert(false);
}
return "";
}
size_t jsonutil_double_to_string(const double val, char *double_to_string_buf, size_t len) {
// It's safe to write a double value into double_to_string_buf, because the converted string will
// never exceed length of 1024.
ValkeyModule_Assert(len == BUF_SIZE_DOUBLE_JSON);
return snprintf(double_to_string_buf, len, "%.17g", val);
}
/**
* Convert double to string using the same format as RapidJSON's Writer::WriteDouble does.
*/
size_t jsonutil_double_to_string_rapidjson(const double val, char *double_to_string_buf_rapidjson, size_t len) {
// RapidJSON's Writer::WriteDouble only uses a buffer of 25 bytes.
ValkeyModule_Assert(len == BUF_SIZE_DOUBLE_RAPID_JSON);
char *end = rapidjson::internal::dtoa(val, double_to_string_buf_rapidjson,
rapidjson::Writer<rapidjson::StringBuffer>::kDefaultMaxDecimalPlaces);
*end = '\0';
return end - double_to_string_buf_rapidjson;
}
bool jsonutil_is_int64(const double a) {
int64_t a_l = static_cast<long long>(a);
double b = static_cast<double>(a_l);
return (a <= b && a >= b);
}
JsonUtilCode jsonutil_multiply_double(const double a, const double b, double *res) {
double c = a * b;
// check overflow
if (std::isinf(c)) return JSONUTIL_MULTIPLICATION_OVERFLOW;
*res = c;
return JSONUTIL_SUCCESS;
}
JsonUtilCode jsonutil_multiply_int64(const int64_t a, const int64_t b, int64_t *res) {
if (a == 0 || b == 0) {
*res = 0;
return JSONUTIL_SUCCESS;
}
// Check overflow conditions without performing multiplication
if ((a > 0 && b > 0 && a > INT64_MAX / b) || // Positive * Positive overflow
(a > 0 && b < 0 && b < INT64_MIN / a) || // Positive * Negative overflow
(a < 0 && b > 0 && a < INT64_MIN / b) || // Negative * Positive overflow
(a < 0 && b < 0 && a < INT64_MAX / b)) { // Negative * Negative overflow
return JSONUTIL_MULTIPLICATION_OVERFLOW;
}
// If no overflow, perform the multiplication
*res = a * b;
return JSONUTIL_SUCCESS;
return JSONUTIL_SUCCESS;
}
JsonUtilCode jsonutil_add_double(const double a, const double b, double *res) {
double c = a + b;
// check overflow
if (std::isinf(c)) return JSONUTIL_ADDITION_OVERFLOW;
*res = c;
return JSONUTIL_SUCCESS;
}
JsonUtilCode jsonutil_add_int64(const int64_t a, const int64_t b, int64_t *res) {
if (a >= 0) {
if (b > INT64_MAX - a) return JSONUTIL_ADDITION_OVERFLOW;
} else {
if (b < INT64_MIN - a) return JSONUTIL_ADDITION_OVERFLOW;
}
*res = a + b;
return JSONUTIL_SUCCESS;
}
bool jsonutil_is_root_path(const char *json_path) {
return !strcmp(json_path, ".") || !strcmp(json_path, "$");
}

122
src/json/util.h Normal file
View File

@ -0,0 +1,122 @@
/**
* This is the utility module, containing shared utility and helper code.
*
* Coding Conventions & Best Practices:
* 1. Every public interface method declared in this file should be prefixed with "jsonutil_".
* 2. Generally speaking, interface methods should not have Valkey module types such as ValkeyModuleCtx
* or ValkeyModuleString, because that would make unit tests hard to write, unless gmock classes
* have been developed.
*/
#ifndef VALKEYJSONMODULE_JSON_UTIL_H_
#define VALKEYJSONMODULE_JSON_UTIL_H_
#include <cmath>
extern "C" {
#define VALKEYMODULE_EXPERIMENTAL_API
#include <./include/valkeymodule.h>
}
typedef enum {
JSONUTIL_SUCCESS = 0,
JSONUTIL_WRONG_NUM_ARGS,
JSONUTIL_JSON_PARSE_ERROR,
JSONUTIL_NX_XX_CONDITION_NOT_SATISFIED,
JSONUTIL_NX_XX_SHOULD_BE_MUTUALLY_EXCLUSIVE,
JSONUTIL_INVALID_JSON_PATH,
JSONUTIL_INVALID_USE_OF_WILDCARD,
JSONUTIL_INVALID_MEMBER_NAME,
JSONUTIL_INVALID_NUMBER,
JSONUTIL_INVALID_IDENTIFIER,
JSONUTIL_INVALID_DOT_SEQUENCE,
JSONUTIL_EMPTY_EXPR_TOKEN,
JSONUTIL_ARRAY_INDEX_NOT_NUMBER,
JSONUTIL_STEP_CANNOT_NOT_BE_ZERO,
JSONUTIL_JSON_PATH_NOT_EXIST,
JSONUTIL_PARENT_ELEMENT_NOT_EXIST,
JSONUTIL_DOCUMENT_KEY_NOT_FOUND,
JSONUTIL_NOT_A_DOCUMENT_KEY,
JSONUTIL_FAILED_TO_DELETE_VALUE,
JSONUTIL_JSON_ELEMENT_NOT_NUMBER,
JSONUTIL_JSON_ELEMENT_NOT_BOOL,
JSONUTIL_JSON_ELEMENT_NOT_STRING,
JSONUTIL_JSON_ELEMENT_NOT_OBJECT,
JSONUTIL_JSON_ELEMENT_NOT_ARRAY,
JSONUTIL_VALUE_NOT_NUMBER,
JSONUTIL_VALUE_NOT_STRING,
JSONUTIL_VALUE_NOT_INTEGER,
JSONUTIL_PATH_SHOULD_BE_AT_THE_END,
JSONUTIL_COMMAND_SYNTAX_ERROR,
JSONUTIL_MULTIPLICATION_OVERFLOW,
JSONUTIL_ADDITION_OVERFLOW,
JSONUTIL_EMPTY_JSON_OBJECT,
JSONUTIL_EMPTY_JSON_ARRAY,
JSONUTIL_INDEX_OUT_OF_ARRAY_BOUNDARIES,
JSONUTIL_UNKNOWN_SUBCOMMAND,
JSONUTIL_FAILED_TO_CREATE_THREAD_SPECIFIC_DATA_KEY,
JSONUTIL_DOCUMENT_SIZE_LIMIT_EXCEEDED,
JSONUTIL_DOCUMENT_PATH_LIMIT_EXCEEDED,
JSONUTIL_PARSER_RECURSION_DEPTH_LIMIT_EXCEEDED,
JSONUTIL_RECURSIVE_DESCENT_TOKEN_LIMIT_EXCEEDED,
JSONUTIL_QUERY_STRING_SIZE_LIMIT_EXCEEDED,
JSONUTIL_CANNOT_INSERT_MEMBER_INTO_NON_OBJECT_VALUE,
JSONUTIL_INVALID_RDB_FORMAT,
JSONUTIL_DOLLAR_CANNOT_APPLY_TO_NON_ROOT,
JSONUTIL_LAST
} JsonUtilCode;
typedef struct {
const char *newline;
const char *space;
const char *indent;
} PrintFormat;
/* Enums for buffer sizes used in conversion of double to json or double to rapidjson */
enum { BUF_SIZE_DOUBLE_JSON = 32, BUF_SIZE_DOUBLE_RAPID_JSON = 25};
/* Get message for a given code. */
const char *jsonutil_code_to_message(JsonUtilCode code);
/* Convert a double value to string. This method is used to help serializing numbers to strings.
* Trailing zeros will be removed. For example, 135.250000 will be converted to string 135.25.
*/
size_t jsonutil_double_to_string(const double val, char *double_to_string_buf, size_t len);
/**
* Convert double to string using the same format as RapidJSON's Writer::WriteDouble does.
*/
size_t jsonutil_double_to_string_rapidjson(const double val, char* double_to_string_buf_rapidjson, size_t len);
/* Check if a double value is int64.
* If the given double does not equal an integer (int64), return false.
* If the given double is out of range of int64, return false.
*/
bool jsonutil_is_int64(const double a);
/* Multiple two double numbers with overflow check.
* @param res - OUTPUT parameter, *res stores the result of multiplication.
* @return JSONUTIL_SUCCESS if successful, JSONUTIL_MULTIPLICATION_OVERFLOW if the result overflows.
*/
JsonUtilCode jsonutil_multiply_double(const double a, const double b, double *res);
/* Multiple two int64 numbers with overflow check.
* @param res - OUTPUT parameter, *res stores the result of multiplication.
* @return JSONUTIL_SUCCESS if successful, JSONUTIL_MULTIPLICATION_OVERFLOW if the result overflows.
*/
JsonUtilCode jsonutil_multiply_int64(const int64_t a, const int64_t b, int64_t *res);
/* Add two double numbers with overflow check.
* @param res - OUTPUT parameter, *res stores the result of addition.
* @return JSONUTIL_SUCCESS if successful, JSONUTIL_ADDITION_OVERFLOW if the result overflows.
*/
JsonUtilCode jsonutil_add_double(const double a, const double b, double *res);
/* Add two int64 numbers with overflow check.
* @param res - OUTPUT parameter, *res stores the result of addition.
* @return JSONUTIL_SUCCESS if successful, JSONUTIL_ADDITION_OVERFLOW if the result overflows.
*/
JsonUtilCode jsonutil_add_int64(const int64_t a, const int64_t b, int64_t *res);
bool jsonutil_is_root_path(const char *json_path);
#endif // VALKEYJSONMODULE_JSON_UTIL_H_

1
src/rapidjson/CPPLINT.cfg Normal file
View File

@ -0,0 +1 @@
exclude_files=.*

14
src/rapidjson/README.md Normal file
View File

@ -0,0 +1,14 @@
# RapidJSON Source Code
* The original RapidJSON Source Code is cloned at build time using CMAKELISTS
* Last commit on the master branch: 0d4517f15a8d7167ba9ae67f3f22a559ca841e3b, 2021-10-31 11:07:57
# Modifications
We made a few changes to the RapidJSON source code. Before the changes are pushed to the open source,
we have to include a private copy of the file. Modified RapidJSON code is under src/rapidjson.
## document.h`
We need to modify RapidJSON's document.h to support JSON depth limit.
### reader.h
Modified reader.h to only generate integers in int64 range.

3599
src/rapidjson/document.h Normal file
View File

File diff suppressed because it is too large Load Diff

57
src/rapidjson/license.txt Normal file
View File

@ -0,0 +1,57 @@
Tencent is pleased to support the open source community by making RapidJSON available.
Copyright (C) 2015 THL A29 Limited, a Tencent company, and Milo Yip. All rights reserved.
If you have downloaded a copy of the RapidJSON binary from Tencent, please note that the RapidJSON binary is licensed under the MIT License.
If you have downloaded a copy of the RapidJSON source code from Tencent, please note that RapidJSON source code is licensed under the MIT License, except for the third-party components listed below which are subject to different license terms. Your integration of RapidJSON into your own projects may require compliance with the MIT License, as well as the other licenses applicable to the third-party components included within RapidJSON. To avoid the problematic JSON license in your own projects, it's sufficient to exclude the bin/jsonchecker/ directory, as it's the only code under the JSON license.
A copy of the MIT License is included in this file.
Other dependencies and licenses:
Open Source Software Licensed Under the BSD License:
--------------------------------------------------------------------
The msinttypes r29
Copyright (c) 2006-2013 Alexander Chemeris
All rights reserved.
Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
* Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
* Neither the name of copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS AND CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
Open Source Software Licensed Under the JSON License:
--------------------------------------------------------------------
json.org
Copyright (c) 2002 JSON.org
All Rights Reserved.
JSON_checker
Copyright (c) 2002 JSON.org
All Rights Reserved.
Terms of the JSON License:
---------------------------------------------------
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
The Software shall be used for Good, not Evil.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
Terms of the MIT License:
--------------------------------------------------------------------
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

393
src/rapidjson/prettywriter.h Normal file
View File

@ -0,0 +1,393 @@
// Tencent is pleased to support the open source community by making RapidJSON available.
//
// Copyright (C) 2015 THL A29 Limited, a Tencent company, and Milo Yip.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
#ifndef RAPIDJSON_PRETTYWRITER_H_
#define RAPIDJSON_PRETTYWRITER_H_
#include <rapidjson/writer.h>
#include "json/json.h"
#ifdef __GNUC__
RAPIDJSON_DIAG_PUSH
RAPIDJSON_DIAG_OFF(effc++)
#endif
#if defined(__clang__)
RAPIDJSON_DIAG_PUSH
RAPIDJSON_DIAG_OFF(c++98-compat)
#endif
RAPIDJSON_NAMESPACE_BEGIN
//! Combination of PrettyWriter format flags.
/*! \see PrettyWriter::SetFormatOptions
*/
enum PrettyFormatOptions {
kFormatDefault = 0, //!< Default pretty formatting.
kFormatSingleLineArray = 1 //!< Format arrays on a single line.
};
//! Writer with indentation and spacing.
/*!
\tparam OutputStream Type of output os.
\tparam SourceEncoding Encoding of source string.
\tparam TargetEncoding Encoding of output stream.
\tparam StackAllocator Type of allocator for allocating memory of stack.
*/
template<typename OutputStream, typename SourceEncoding = UTF8<>, typename TargetEncoding = UTF8<>, typename StackAllocator = CrtAllocator, unsigned writeFlags = kWriteDefaultFlags>
class PrettyWriter : public Writer<OutputStream, SourceEncoding, TargetEncoding, StackAllocator, writeFlags> {
public:
typedef Writer<OutputStream, SourceEncoding, TargetEncoding, StackAllocator, writeFlags> Base;
typedef typename Base::Ch Ch;
//! Constructor
/*! \param os Output stream.
\param allocator User supplied allocator. If it is null, it will create a private one.
\param levelDepth Initial capacity of stack.
*/
explicit PrettyWriter(OutputStream& os, StackAllocator* allocator = 0, size_t levelDepth = Base::kDefaultLevelDepth) :
Base(os, allocator, levelDepth), formatOptions_(kFormatDefault), initialLevel(0), curDepth(0), maxDepth(0) {}
explicit PrettyWriter(StackAllocator* allocator = 0, size_t levelDepth = Base::kDefaultLevelDepth) :
Base(allocator, levelDepth), formatOptions_(kFormatDefault), initialLevel(0), curDepth(0), maxDepth(0) {}
#if RAPIDJSON_HAS_CXX11_RVALUE_REFS
PrettyWriter(PrettyWriter&& rhs) :
Base(std::forward<PrettyWriter>(rhs)), formatOptions_(rhs.formatOptions_),
newline_(rhs.newline_), indent_(rhs.indent_), space_(rhs.space_), initialLevel(rhs.initialLevel_) {}
#endif
//! Set pretty writer formatting options.
/*! \param options Formatting options.
*/
PrettyWriter& SetFormatOptions(PrettyFormatOptions options) {
formatOptions_ = options;
return *this;
}
PrettyWriter& SetNewline(const std::string_view &newline) {
newline_ = newline;
return *this;
}
PrettyWriter& SetIndent(const std::string_view &indent) {
indent_ = indent;
return *this;
}
PrettyWriter& SetSpace(const std::string_view &space) {
space_ = space;
return *this;
}
PrettyWriter& SetInitialLevel(size_t il) {
initialLevel = il;
return *this;
}
/*! @name Implementation of Handler
\see Handler
*/
//@{
bool Null() { PrettyPrefix(kNullType); return Base::EndValue(Base::WriteNull()); }
bool Bool(bool b) { PrettyPrefix(b ? kTrueType : kFalseType); return Base::EndValue(Base::WriteBool(b)); }
bool Int(int i) { PrettyPrefix(kNumberType); return Base::EndValue(Base::WriteInt(i)); }
bool Uint(unsigned u) { PrettyPrefix(kNumberType); return Base::EndValue(Base::WriteUint(u)); }
bool Int64(int64_t i64) { PrettyPrefix(kNumberType); return Base::EndValue(Base::WriteInt64(i64)); }
bool Uint64(uint64_t u64) { PrettyPrefix(kNumberType); return Base::EndValue(Base::WriteUint64(u64)); }
bool Double(double d) { PrettyPrefix(kNumberType); return Base::EndValue(Base::WriteDouble(d)); }
bool RawNumber(const Ch* str, SizeType length, bool copy = false) {
RAPIDJSON_ASSERT(str != 0);
(void)copy;
PrettyPrefix(kNumberType);
return Base::EndValue(Base::WriteDouble(str, length));
}
bool String(const Ch* str, SizeType length, bool copy = false) {
RAPIDJSON_ASSERT(str != 0);
(void)copy;
PrettyPrefix(kStringType);
return Base::EndValue(Base::WriteString(str, length));
}
#if RAPIDJSON_HAS_STDSTRING
bool String(const std::basic_string<Ch>& str) {
return String(str.data(), SizeType(str.size()));
}
#endif
size_t GetMaxDepth() {
return maxDepth;
}
bool StartObject() {
IncrDepth();
PrettyPrefix(kObjectType);
new (Base::level_stack_.template Push<typename Base::Level>()) typename Base::Level(false);
return Base::WriteStartObject();
}
bool Key(const Ch* str, SizeType length, bool copy = false) { return String(str, length, copy); }
#if RAPIDJSON_HAS_STDSTRING
bool Key(const std::basic_string<Ch>& str) {
return Key(str.data(), SizeType(str.size()));
}
#endif
bool EndObject(SizeType memberCount = 0) {
(void)memberCount;
RAPIDJSON_ASSERT(Base::level_stack_.GetSize() >= sizeof(typename Base::Level)); // not inside an Object
RAPIDJSON_ASSERT(!Base::level_stack_.template Top<typename Base::Level>()->inArray); // currently inside an Array, not Object
RAPIDJSON_ASSERT(0 == Base::level_stack_.template Top<typename Base::Level>()->valueCount % 2); // Object has a Key without a Value
bool empty = Base::level_stack_.template Pop<typename Base::Level>(1)->valueCount == 0;
if (!empty) {
WriteNewline();
WriteIndent();
}
bool ret = Base::EndValue(Base::WriteEndObject());
(void)ret;
RAPIDJSON_ASSERT(ret == true);
if (Base::level_stack_.Empty()) // end of json text
Base::Flush();
DecrDepth();
return true;
}
bool StartArray() {
IncrDepth();
PrettyPrefix(kArrayType);
new (Base::level_stack_.template Push<typename Base::Level>()) typename Base::Level(true);
return Base::WriteStartArray();
}
bool EndArray(SizeType memberCount = 0) {
(void)memberCount;
RAPIDJSON_ASSERT(Base::level_stack_.GetSize() >= sizeof(typename Base::Level));
RAPIDJSON_ASSERT(Base::level_stack_.template Top<typename Base::Level>()->inArray);
bool empty = Base::level_stack_.template Pop<typename Base::Level>(1)->valueCount == 0;
if (!empty && !(formatOptions_ & kFormatSingleLineArray)) {
WriteNewline();
WriteIndent();
}
bool ret = Base::EndValue(Base::WriteEndArray());
(void)ret;
RAPIDJSON_ASSERT(ret == true);
if (Base::level_stack_.Empty()) // end of json text
Base::Flush();
DecrDepth();
return true;
}
//@}
/*! @name Convenience extensions */
//@{
//! Simpler but slower overload.
bool String(const Ch* str) { return String(str, internal::StrLen(str)); }
bool Key(const Ch* str) { return Key(str, internal::StrLen(str)); }
//@}
//! Write a raw JSON value.
/*!
For user to write a stringified JSON as a value.
\param json A well-formed JSON value. It should not contain null character within [0, length - 1] range.
\param length Length of the json.
\param type Type of the root of json.
\note When using PrettyWriter::RawValue(), the result json may not be indented correctly.
*/
bool RawValue(const Ch* json, size_t length, Type type) {
RAPIDJSON_ASSERT(json != 0);
PrettyPrefix(type);
return Base::EndValue(Base::WriteRawValue(json, length));
}
protected:
void PrettyPrefix(Type type) {
(void)type;
if (Base::level_stack_.GetSize() != 0) { // this value is not at root
typename Base::Level* level = Base::level_stack_.template Top<typename Base::Level>();
if (level->inArray) {
if (level->valueCount > 0) {
Base::os_->Put(','); // add comma if it is not the first element in array
if (formatOptions_ & kFormatSingleLineArray)
WriteSpace();
}
if (!(formatOptions_ & kFormatSingleLineArray)) {
WriteNewline();
WriteIndent();
}
}
else { // in object
if (level->valueCount > 0) {
if (level->valueCount % 2 == 0) {
Base::os_->Put(',');
WriteNewline();
}
else {
Base::os_->Put(':');
WriteSpace();
}
}
else
WriteNewline();
if (level->valueCount % 2 == 0)
WriteIndent();
}
if (!level->inArray && level->valueCount % 2 == 0)
RAPIDJSON_ASSERT(type == kStringType); // if it's in object, then even number should be a name
level->valueCount++;
}
else {
RAPIDJSON_ASSERT(!Base::hasRoot_); // Should only has one and only one root.
Base::hasRoot_ = true;
}
}
void WriteStringView(const std::string_view& v) {
if (!v.empty()) {
size_t sz = v.size();
char *buf = Base::os_->Push(sz);
v.copy(buf, sz);
}
}
void WriteString(const char *ptr, size_t len, bool noescape) {
if (noescape) {
char *p = Base::os_->Push(len + 2);
p[0] = '"';
std::memcpy(p + 1, ptr, len);
p[len + 1] = '"';
} else {
Base::WriteString(ptr, len);
}
}
void WriteNewline() { WriteStringView(newline_); }
void WriteSpace() { WriteStringView(space_); }
void WriteIndent() {
size_t count = initialLevel + (Base::level_stack_.GetSize() / sizeof(typename Base::Level));
for (size_t i = 0; i < count; ++i) WriteStringView(indent_);
}
public:
//
// Accelerated write when there's definitely no format
//
template<typename JValue>
void FastWrite(JValue &value, size_t *max_depth) {
*max_depth = 0;
FastWrite_internal(value, 0, max_depth);
}
PrettyFormatOptions formatOptions_;
std::string_view newline_;
std::string_view indent_;
std::string_view space_;
size_t initialLevel;
private:
// Prohibit copy constructor & assignment operator.
PrettyWriter(const PrettyWriter&);
PrettyWriter& operator=(const PrettyWriter&);
size_t curDepth;
size_t maxDepth;
void IncrDepth() {
curDepth++;
if (curDepth > maxDepth) maxDepth = curDepth;
}
void DecrDepth() {
RAPIDJSON_ASSERT(curDepth > 0);
curDepth--;
}
template<typename JValue>
void FastWrite_internal(JValue &value, const size_t level, size_t *max_depth) {
if (level > *max_depth) *max_depth = level;
bool firstElement;
switch (value.GetType()) {
case kStringType:
WriteString(value.GetString(), value.GetStringLength(), value.IsNoescape());
break;
case kNullType:
Base::WriteNull();
break;
case kFalseType:
Base::WriteBool(false);
break;
case kTrueType:
Base::WriteBool(true);
break;
case kObjectType:
Base::os_->Put('{');
firstElement = true;
for (typename JValue::ConstMemberIterator m = value.MemberBegin(); m != value.MemberEnd(); ++m) {
if (!firstElement) {
Base::os_->Put(',');
} else {
firstElement = false;
}
WriteString(m->name.GetString(), m->name.GetStringLength(), m->name.IsNoescape());
Base::os_->Put(':');
FastWrite_internal(m->value, level + 1, max_depth);
}
Base::os_->Put('}');
break;
case kArrayType:
Base::os_->Put('[');
firstElement = true;
for (typename JValue::ConstValueIterator v = value.Begin(); v != value.End(); ++v) {
if (!firstElement) {
Base::os_->Put(',');
} else {
firstElement = false;
}
FastWrite_internal(*v, level + 1, max_depth);
}
Base::os_->Put(']');
break;
default:
RAPIDJSON_ASSERT(value.GetType() == kNumberType);
if (value.IsDouble()) {
Base::WriteDouble(value.GetDoubleString(), value.GetDoubleStringLength());
}
else if (value.IsInt()) Base::WriteInt(value.GetInt());
else if (value.IsUint()) Base::WriteUint(value.GetUint());
else if (value.IsInt64()) Base::WriteInt64(value.GetInt64());
else Base::WriteUint64(value.GetUint64());
break;
}
}
};
RAPIDJSON_NAMESPACE_END
#if defined(__clang__)
RAPIDJSON_DIAG_POP
#endif
#ifdef __GNUC__
RAPIDJSON_DIAG_POP
#endif
#endif // RAPIDJSON_RAPIDJSON_H_

2281
src/rapidjson/reader.h Normal file
View File

File diff suppressed because it is too large Load Diff

119
src/rapidjson/stringbuffer.h Normal file
View File

@ -0,0 +1,119 @@
// Tencent is pleased to support the open source community by making RapidJSON available.
//
// Copyright (C) 2015 THL A29 Limited, a Tencent company, and Milo Yip.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
#ifndef RAPIDJSON_STRINGBUFFER_H_
#define RAPIDJSON_STRINGBUFFER_H_
#include <rapidjson/stream.h>
#include <rapidjson/internal/stack.h>
#if RAPIDJSON_HAS_CXX11_RVALUE_REFS
#include <utility> // std::move
#endif
#if defined(__clang__)
RAPIDJSON_DIAG_PUSH
RAPIDJSON_DIAG_OFF(c++98-compat)
#endif
RAPIDJSON_NAMESPACE_BEGIN
//! Represents an in-memory output stream.
/*!
\tparam Encoding Encoding of the stream.
\tparam Allocator type for allocating memory buffer.
\note implements Stream concept
*/
template <typename Encoding, typename Allocator = CrtAllocator>
class GenericStringBuffer {
public:
typedef typename Encoding::Ch Ch;
GenericStringBuffer(Allocator* allocator = 0, size_t capacity = kDefaultCapacity) : stack_(allocator, capacity) {}
#if RAPIDJSON_HAS_CXX11_RVALUE_REFS
GenericStringBuffer(GenericStringBuffer&& rhs) : stack_(std::move(rhs.stack_)) {}
GenericStringBuffer& operator=(GenericStringBuffer&& rhs) {
if (&rhs != this)
stack_ = std::move(rhs.stack_);
return *this;
}
#endif
void Put(Ch c) { *stack_.template Push<Ch>() = c; }
void PutUnsafe(Ch c) { *stack_.template PushUnsafe<Ch>() = c; }
void Flush() {}
void Clear() { stack_.Clear(); }
void ShrinkToFit() {
// Push and pop a null terminator. This is safe.
*stack_.template Push<Ch>() = '\0';
stack_.ShrinkToFit();
stack_.template Pop<Ch>(1);
}
void Reserve(size_t count) { stack_.template Reserve<Ch>(count); }
Ch* Push(size_t count) { return stack_.template Push<Ch>(count); }
Ch* PushUnsafe(size_t count) { return stack_.template PushUnsafe<Ch>(count); }
void Pop(size_t count) { stack_.template Pop<Ch>(count); }
const Ch* GetString() const {
// Push and pop a null terminator. This is safe.
*stack_.template Push<Ch>() = '\0';
stack_.template Pop<Ch>(1);
return stack_.template Bottom<Ch>();
}
//! Get the size of string in bytes in the string buffer.
size_t GetSize() const { return stack_.GetSize(); }
//! Get the length of string in Ch in the string buffer.
size_t GetLength() const { return stack_.GetSize() / sizeof(Ch); }
static const size_t kDefaultCapacity = 256;
mutable internal::Stack<Allocator> stack_;
private:
// Prohibit copy constructor & assignment operator.
GenericStringBuffer(const GenericStringBuffer&);
GenericStringBuffer& operator=(const GenericStringBuffer&);
};
//! String buffer with UTF8 encoding
typedef GenericStringBuffer<UTF8<> > StringBuffer;
template<typename Encoding, typename Allocator>
inline void PutReserve(GenericStringBuffer<Encoding, Allocator>& stream, size_t count) {
stream.Reserve(count);
}
template<typename Encoding, typename Allocator>
inline void PutUnsafe(GenericStringBuffer<Encoding, Allocator>& stream, typename Encoding::Ch c) {
stream.PutUnsafe(c);
}
//! Implement specialized version of PutN() with memset() for better performance.
template<>
inline void PutN(GenericStringBuffer<UTF8<> >& stream, char c, size_t n) {
std::memset(stream.stack_.Push<char>(n), c, n * sizeof(c));
}
RAPIDJSON_NAMESPACE_END
#if defined(__clang__)
RAPIDJSON_DIAG_POP
#endif
#endif // RAPIDJSON_STRINGBUFFER_H_

730
src/rapidjson/writer.h Normal file
View File

@ -0,0 +1,730 @@
// Tencent is pleased to support the open source community by making RapidJSON available.
//
// Copyright (C) 2015 THL A29 Limited, a Tencent company, and Milo Yip.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
#ifndef RAPIDJSON_WRITER_H_
#define RAPIDJSON_WRITER_H_
#include <rapidjson/stream.h>
#include <rapidjson/internal/clzll.h>
#include <rapidjson/internal/meta.h>
#include <rapidjson/internal/stack.h>
#include <rapidjson/internal/strfunc.h>
#include <rapidjson/internal/dtoa.h>
#include <rapidjson/internal/itoa.h>
#include "stringbuffer.h"
#include <iostream>
#include <new> // placement new
#if defined(RAPIDJSON_SIMD) && defined(_MSC_VER)
#include <intrin.h>
#pragma intrinsic(_BitScanForward)
#endif
#ifdef RAPIDJSON_SSE42
#include <nmmintrin.h>
#elif defined(RAPIDJSON_SSE2)
#include <emmintrin.h>
#elif defined(RAPIDJSON_NEON)
#include <arm_neon.h>
#endif
#ifdef __clang__
RAPIDJSON_DIAG_PUSH
RAPIDJSON_DIAG_OFF(padded)
RAPIDJSON_DIAG_OFF(unreachable-code)
RAPIDJSON_DIAG_OFF(c++98-compat)
#elif defined(_MSC_VER)
RAPIDJSON_DIAG_PUSH
RAPIDJSON_DIAG_OFF(4127) // conditional expression is constant
#endif
RAPIDJSON_NAMESPACE_BEGIN
///////////////////////////////////////////////////////////////////////////////
// WriteFlag
/*! \def RAPIDJSON_WRITE_DEFAULT_FLAGS
\ingroup RAPIDJSON_CONFIG
\brief User-defined kWriteDefaultFlags definition.
User can define this as any \c WriteFlag combinations.
*/
#ifndef RAPIDJSON_WRITE_DEFAULT_FLAGS
#define RAPIDJSON_WRITE_DEFAULT_FLAGS kWriteNoFlags
#endif
template<class StringBuffer> void Put(StringBuffer& os, typename StringBuffer::Ch c) { os.Put(c); }
//! Combination of writeFlags
enum WriteFlag {
kWriteNoFlags = 0, //!< No flags are set.
kWriteValidateEncodingFlag = 1, //!< Validate encoding of JSON strings.
kWriteNanAndInfFlag = 2, //!< Allow writing of Infinity, -Infinity and NaN.
kWriteDefaultFlags = RAPIDJSON_WRITE_DEFAULT_FLAGS //!< Default write flags. Can be customized by defining RAPIDJSON_WRITE_DEFAULT_FLAGS
};
//! JSON writer
/*! Writer implements the concept Handler.
It generates JSON text by events to an output os.
User may programmatically calls the functions of a writer to generate JSON text.
On the other side, a writer can also be passed to objects that generates events,
for example Reader::Parse() and Document::Accept().
\tparam OutputStream Type of output stream.
\tparam SourceEncoding Encoding of source string.
\tparam TargetEncoding Encoding of output stream.
\tparam StackAllocator Type of allocator for allocating memory of stack.
\note implements Handler concept
*/
template<typename OutputStream, typename SourceEncoding = UTF8<>, typename TargetEncoding = UTF8<>, typename StackAllocator = CrtAllocator, unsigned writeFlags = kWriteDefaultFlags>
class Writer {
public:
typedef typename SourceEncoding::Ch Ch;
static const int kDefaultMaxDecimalPlaces = 324;
//! Constructor
/*! \param os Output stream.
\param stackAllocator User supplied allocator. If it is null, it will create a private one.
\param levelDepth Initial capacity of stack.
*/
explicit
Writer(OutputStream& os, StackAllocator* stackAllocator = 0, size_t levelDepth = kDefaultLevelDepth) :
os_(&os), level_stack_(stackAllocator, levelDepth * sizeof(Level)), maxDecimalPlaces_(kDefaultMaxDecimalPlaces), hasRoot_(false) {}
explicit
Writer(StackAllocator* allocator = 0, size_t levelDepth = kDefaultLevelDepth) :
os_(0), level_stack_(allocator, levelDepth * sizeof(Level)), maxDecimalPlaces_(kDefaultMaxDecimalPlaces), hasRoot_(false) {}
#if RAPIDJSON_HAS_CXX11_RVALUE_REFS
Writer(Writer&& rhs) :
os_(rhs.os_), level_stack_(std::move(rhs.level_stack_)), maxDecimalPlaces_(rhs.maxDecimalPlaces_), hasRoot_(rhs.hasRoot_) {
rhs.os_ = 0;
}
#endif
//! Reset the writer with a new stream.
/*!
This function reset the writer with a new stream and default settings,
in order to make a Writer object reusable for output multiple JSONs.
\param os New output stream.
\code
Writer<OutputStream> writer(os1);
writer.StartObject();
// ...
writer.EndObject();
writer.Reset(os2);
writer.StartObject();
// ...
writer.EndObject();
\endcode
*/
void Reset(OutputStream& os) {
os_ = &os;
hasRoot_ = false;
level_stack_.Clear();
}
//! Checks whether the output is a complete JSON.
/*!
A complete JSON has a complete root object or array.
*/
bool IsComplete() const {
return hasRoot_ && level_stack_.Empty();
}
int GetMaxDecimalPlaces() const {
return maxDecimalPlaces_;
}
//! Sets the maximum number of decimal places for double output.
/*!
This setting truncates the output with specified number of decimal places.
For example,
\code
writer.SetMaxDecimalPlaces(3);
writer.StartArray();
writer.Double(0.12345); // "0.123"
writer.Double(0.0001); // "0.0"
writer.Double(1.234567890123456e30); // "1.234567890123456e30" (do not truncate significand for positive exponent)
writer.Double(1.23e-4); // "0.0" (do truncate significand for negative exponent)
writer.EndArray();
\endcode
The default setting does not truncate any decimal places. You can restore to this setting by calling
\code
writer.SetMaxDecimalPlaces(Writer::kDefaultMaxDecimalPlaces);
\endcode
*/
void SetMaxDecimalPlaces(int maxDecimalPlaces) {
maxDecimalPlaces_ = maxDecimalPlaces;
}
/*!@name Implementation of Handler
\see Handler
*/
//@{
bool Null() { Prefix(kNullType); return EndValue(WriteNull()); }
bool Bool(bool b) { Prefix(b ? kTrueType : kFalseType); return EndValue(WriteBool(b)); }
bool Int(int i) { Prefix(kNumberType); return EndValue(WriteInt(i)); }
bool Uint(unsigned u) { Prefix(kNumberType); return EndValue(WriteUint(u)); }
bool Int64(int64_t i64) { Prefix(kNumberType); return EndValue(WriteInt64(i64)); }
bool Uint64(uint64_t u64) { Prefix(kNumberType); return EndValue(WriteUint64(u64)); }
//! Writes the given \c double value to the stream
/*!
\param d The value to be written.
\return Whether it is succeed.
*/
bool Double(double d) { Prefix(kNumberType); return EndValue(WriteDouble(d)); }
bool RawNumber(const Ch* str, SizeType length, bool copy = false) {
RAPIDJSON_ASSERT(str != 0);
(void)copy;
Prefix(kNumberType);
return EndValue(WriteDouble(str, length));
}
bool String(const Ch* str, SizeType length, bool copy = false) {
RAPIDJSON_ASSERT(str != 0);
(void)copy;
Prefix(kStringType);
return EndValue(WriteString(str, length));
}
#if RAPIDJSON_HAS_STDSTRING
bool String(const std::basic_string<Ch>& str) {
return String(str.data(), SizeType(str.size()));
}
#endif
bool StartObject() {
Prefix(kObjectType);
new (level_stack_.template Push<Level>()) Level(false);
return WriteStartObject();
}
bool Key(const Ch* str, SizeType length, bool copy = false) { return String(str, length, copy); }
#if RAPIDJSON_HAS_STDSTRING
bool Key(const std::basic_string<Ch>& str)
{
return Key(str.data(), SizeType(str.size()));
}
#endif
bool EndObject(SizeType memberCount = 0) {
(void)memberCount;
RAPIDJSON_ASSERT(level_stack_.GetSize() >= sizeof(Level)); // not inside an Object
RAPIDJSON_ASSERT(!level_stack_.template Top<Level>()->inArray); // currently inside an Array, not Object
RAPIDJSON_ASSERT(0 == level_stack_.template Top<Level>()->valueCount % 2); // Object has a Key without a Value
level_stack_.template Pop<Level>(1);
return EndValue(WriteEndObject());
}
bool StartArray() {
Prefix(kArrayType);
new (level_stack_.template Push<Level>()) Level(true);
return WriteStartArray();
}
bool EndArray(SizeType elementCount = 0) {
(void)elementCount;
RAPIDJSON_ASSERT(level_stack_.GetSize() >= sizeof(Level));
RAPIDJSON_ASSERT(level_stack_.template Top<Level>()->inArray);
level_stack_.template Pop<Level>(1);
return EndValue(WriteEndArray());
}
//@}
/*! @name Convenience extensions */
//@{
//! Simpler but slower overload.
bool String(const Ch* const& str) { return String(str, internal::StrLen(str)); }
bool Key(const Ch* const& str) { return Key(str, internal::StrLen(str)); }
//@}
//! Write a raw JSON value.
/*!
For user to write a stringified JSON as a value.
\param json A well-formed JSON value. It should not contain null character within [0, length - 1] range.
\param length Length of the json.
\param type Type of the root of json.
*/
bool RawValue(const Ch* json, size_t length, Type type) {
RAPIDJSON_ASSERT(json != 0);
Prefix(type);
return EndValue(WriteRawValue(json, length));
}
//! Flush the output stream.
/*!
Allows the user to flush the output stream immediately.
*/
void Flush() {
os_->Flush();
}
static const size_t kDefaultLevelDepth = 32;
protected:
//! Information for each nested level
struct Level {
Level(bool inArray_) : valueCount(0), inArray(inArray_) {}
size_t valueCount; //!< number of values in this level
bool inArray; //!< true if in array, otherwise in object
};
bool WriteNull() {
PutReserve(*os_, 4);
PutUnsafe(*os_, 'n'); PutUnsafe(*os_, 'u'); PutUnsafe(*os_, 'l'); PutUnsafe(*os_, 'l'); return true;
}
bool WriteBool(bool b) {
if (b) {
PutReserve(*os_, 4);
PutUnsafe(*os_, 't'); PutUnsafe(*os_, 'r'); PutUnsafe(*os_, 'u'); PutUnsafe(*os_, 'e');
}
else {
PutReserve(*os_, 5);
PutUnsafe(*os_, 'f'); PutUnsafe(*os_, 'a'); PutUnsafe(*os_, 'l'); PutUnsafe(*os_, 's'); PutUnsafe(*os_, 'e');
}
return true;
}
bool WriteInt(int i) {
char buffer[11];
const char* end = internal::i32toa(i, buffer);
PutReserve(*os_, static_cast<size_t>(end - buffer));
for (const char* p = buffer; p != end; ++p)
PutUnsafe(*os_, static_cast<typename OutputStream::Ch>(*p));
return true;
}
bool WriteUint(unsigned u) {
char buffer[10];
const char* end = internal::u32toa(u, buffer);
PutReserve(*os_, static_cast<size_t>(end - buffer));
for (const char* p = buffer; p != end; ++p)
PutUnsafe(*os_, static_cast<typename OutputStream::Ch>(*p));
return true;
}
bool WriteInt64(int64_t i64) {
char buffer[21];
const char* end = internal::i64toa(i64, buffer);
PutReserve(*os_, static_cast<size_t>(end - buffer));
for (const char* p = buffer; p != end; ++p)
PutUnsafe(*os_, static_cast<typename OutputStream::Ch>(*p));
return true;
}
bool WriteUint64(uint64_t u64) {
char buffer[20];
char* end = internal::u64toa(u64, buffer);
PutReserve(*os_, static_cast<size_t>(end - buffer));
for (char* p = buffer; p != end; ++p)
PutUnsafe(*os_, static_cast<typename OutputStream::Ch>(*p));
return true;
}
bool WriteDouble(double d) {
if (internal::Double(d).IsNanOrInf()) {
if (!(writeFlags & kWriteNanAndInfFlag))
return false;
if (internal::Double(d).IsNan()) {
PutReserve(*os_, 3);
PutUnsafe(*os_, 'N'); PutUnsafe(*os_, 'a'); PutUnsafe(*os_, 'N');
return true;
}
if (internal::Double(d).Sign()) {
PutReserve(*os_, 9);
PutUnsafe(*os_, '-');
}
else
PutReserve(*os_, 8);
PutUnsafe(*os_, 'I'); PutUnsafe(*os_, 'n'); PutUnsafe(*os_, 'f');
PutUnsafe(*os_, 'i'); PutUnsafe(*os_, 'n'); PutUnsafe(*os_, 'i'); PutUnsafe(*os_, 't'); PutUnsafe(*os_, 'y');
return true;
}
char buffer[25];
char* end = internal::dtoa(d, buffer, maxDecimalPlaces_);
PutReserve(*os_, static_cast<size_t>(end - buffer));
for (char* p = buffer; p != end; ++p)
PutUnsafe(*os_, static_cast<typename OutputStream::Ch>(*p));
return true;
}
bool WriteDouble(const Ch* str, SizeType length) {
PutReserve(*os_, length);
for (const char* p = str; p!= str+length; ++p)
PutUnsafe(*os_, static_cast<typename OutputStream::Ch>(*p));
return true;
}
bool WriteString(const Ch* str, SizeType length) {
static const typename OutputStream::Ch hexDigits[16] = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F' };
static const char escape[256] = {
#define Z16 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
//0 1 2 3 4 5 6 7 8 9 A B C D E F
'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', 'b', 't', 'n', 'u', 'f', 'r', 'u', 'u', // 00
'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', 'u', // 10
0, 0, '"', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 20
Z16, Z16, // 30~4F
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,'\\', 0, 0, 0, // 50
Z16, Z16, Z16, Z16, Z16, Z16, Z16, Z16, Z16, Z16 // 60~FF
#undef Z16
};
//
// efficiency wants us to PutReserve far in advance
// But that causes memory allocation issues with Chunked Buffer, so don't reserve too far into the future...
//
Put(*os_, '\"');
GenericStringStream<SourceEncoding> is(str);
while (ScanWriteUnescapedString(is, length)) {
const Ch c = is.Peek();
if (!TargetEncoding::supportUnicode && static_cast<unsigned>(c) >= 0x80) {
// Unicode escaping
unsigned codepoint;
if (RAPIDJSON_UNLIKELY(!SourceEncoding::Decode(is, &codepoint)))
return false;
PutReserve(*os_, 12);
PutUnsafe(*os_, '\\');
PutUnsafe(*os_, 'u');
if (codepoint <= 0xD7FF || (codepoint >= 0xE000 && codepoint <= 0xFFFF)) {
PutUnsafe(*os_, hexDigits[(codepoint >> 12) & 15]);
PutUnsafe(*os_, hexDigits[(codepoint >> 8) & 15]);
PutUnsafe(*os_, hexDigits[(codepoint >> 4) & 15]);
PutUnsafe(*os_, hexDigits[(codepoint ) & 15]);
}
else {
RAPIDJSON_ASSERT(codepoint >= 0x010000 && codepoint <= 0x10FFFF);
// Surrogate pair
unsigned s = codepoint - 0x010000;
unsigned lead = (s >> 10) + 0xD800;
unsigned trail = (s & 0x3FF) + 0xDC00;
PutUnsafe(*os_, hexDigits[(lead >> 12) & 15]);
PutUnsafe(*os_, hexDigits[(lead >> 8) & 15]);
PutUnsafe(*os_, hexDigits[(lead >> 4) & 15]);
PutUnsafe(*os_, hexDigits[(lead ) & 15]);
PutUnsafe(*os_, '\\');
PutUnsafe(*os_, 'u');
PutUnsafe(*os_, hexDigits[(trail >> 12) & 15]);
PutUnsafe(*os_, hexDigits[(trail >> 8) & 15]);
PutUnsafe(*os_, hexDigits[(trail >> 4) & 15]);
PutUnsafe(*os_, hexDigits[(trail ) & 15]);
}
}
else if ((sizeof(Ch) == 1 || static_cast<unsigned>(c) < 256) && RAPIDJSON_UNLIKELY(escape[static_cast<unsigned char>(c)])) {
is.Take();
PutReserve(*os_, 6);
PutUnsafe(*os_, '\\');
PutUnsafe(*os_, static_cast<typename OutputStream::Ch>(escape[static_cast<unsigned char>(c)]));
if (escape[static_cast<unsigned char>(c)] == 'u') {
PutUnsafe(*os_, '0');
PutUnsafe(*os_, '0');
PutUnsafe(*os_, hexDigits[static_cast<unsigned char>(c) >> 4]);
PutUnsafe(*os_, hexDigits[static_cast<unsigned char>(c) & 0xF]);
}
}
else {
PutReserve(*os_,16); // I think worst case is only 12, but so what.....
if (RAPIDJSON_UNLIKELY(!(writeFlags & kWriteValidateEncodingFlag ?
Transcoder<SourceEncoding, TargetEncoding>::Validate(is, *os_) :
Transcoder<SourceEncoding, TargetEncoding>::TranscodeUnsafe(is, *os_))))
return false;
}
}
Put(*os_, '\"');
return true;
}
bool ScanWriteUnescapedString(GenericStringStream<SourceEncoding>& is, size_t length) {
return RAPIDJSON_LIKELY(is.Tell() < length);
}
bool WriteStartObject() { os_->Put('{'); return true; }
bool WriteEndObject() { os_->Put('}'); return true; }
bool WriteStartArray() { os_->Put('['); return true; }
bool WriteEndArray() { os_->Put(']'); return true; }
bool WriteRawValue(const Ch* json, size_t length) {
PutReserve(*os_, length);
GenericStringStream<SourceEncoding> is(json);
while (RAPIDJSON_LIKELY(is.Tell() < length)) {
RAPIDJSON_ASSERT(is.Peek() != '\0');
if (RAPIDJSON_UNLIKELY(!(writeFlags & kWriteValidateEncodingFlag ?
Transcoder<SourceEncoding, TargetEncoding>::Validate(is, *os_) :
Transcoder<SourceEncoding, TargetEncoding>::TranscodeUnsafe(is, *os_))))
return false;
}
return true;
}
void Prefix(Type type) {
(void)type;
if (RAPIDJSON_LIKELY(level_stack_.GetSize() != 0)) { // this value is not at root
Level* level = level_stack_.template Top<Level>();
if (level->valueCount > 0) {
if (level->inArray)
os_->Put(','); // add comma if it is not the first element in array
else // in object
os_->Put((level->valueCount % 2 == 0) ? ',' : ':');
}
if (!level->inArray && level->valueCount % 2 == 0)
RAPIDJSON_ASSERT(type == kStringType); // if it's in object, then even number should be a name
level->valueCount++;
}
else {
RAPIDJSON_ASSERT(!hasRoot_); // Should only has one and only one root.
hasRoot_ = true;
}
}
// Flush the value if it is the top level one.
bool EndValue(bool ret) {
if (RAPIDJSON_UNLIKELY(level_stack_.Empty())) // end of json text
Flush();
return ret;
}
OutputStream* os_;
internal::Stack<StackAllocator> level_stack_;
int maxDecimalPlaces_;
bool hasRoot_;
private:
// Prohibit copy constructor & assignment operator.
Writer(const Writer&);
Writer& operator=(const Writer&);
};
// Full specialization for StringStream to prevent memory copying
template<>
inline bool Writer<StringBuffer>::WriteInt(int i) {
char *buffer = os_->Push(11);
const char* end = internal::i32toa(i, buffer);
os_->Pop(static_cast<size_t>(11 - (end - buffer)));
return true;
}
template<>
inline bool Writer<StringBuffer>::WriteUint(unsigned u) {
char *buffer = os_->Push(10);
const char* end = internal::u32toa(u, buffer);
os_->Pop(static_cast<size_t>(10 - (end - buffer)));
return true;
}
template<>
inline bool Writer<StringBuffer>::WriteInt64(int64_t i64) {
char *buffer = os_->Push(21);
const char* end = internal::i64toa(i64, buffer);
os_->Pop(static_cast<size_t>(21 - (end - buffer)));
return true;
}
template<>
inline bool Writer<StringBuffer>::WriteUint64(uint64_t u) {
char *buffer = os_->Push(20);
const char* end = internal::u64toa(u, buffer);
os_->Pop(static_cast<size_t>(20 - (end - buffer)));
return true;
}
template<>
inline bool Writer<StringBuffer>::WriteDouble(double d) {
if (internal::Double(d).IsNanOrInf()) {
// Note: This code path can only be reached if (RAPIDJSON_WRITE_DEFAULT_FLAGS & kWriteNanAndInfFlag).
if (!(kWriteDefaultFlags & kWriteNanAndInfFlag))
return false;
if (internal::Double(d).IsNan()) {
PutReserve(*os_, 3);
PutUnsafe(*os_, 'N'); PutUnsafe(*os_, 'a'); PutUnsafe(*os_, 'N');
return true;
}
if (internal::Double(d).Sign()) {
PutReserve(*os_, 9);
PutUnsafe(*os_, '-');
}
else
PutReserve(*os_, 8);
PutUnsafe(*os_, 'I'); PutUnsafe(*os_, 'n'); PutUnsafe(*os_, 'f');
PutUnsafe(*os_, 'i'); PutUnsafe(*os_, 'n'); PutUnsafe(*os_, 'i'); PutUnsafe(*os_, 't'); PutUnsafe(*os_, 'y');
return true;
}
char *buffer = os_->Push(25);
char* end = internal::dtoa(d, buffer, maxDecimalPlaces_);
os_->Pop(static_cast<size_t>(25 - (end - buffer)));
return true;
}
#if (defined(RAPIDJSON_SSE2) || defined(RAPIDJSON_SSE42)) && !defined(__SANITIZE_ADDRESS__)
template<>
inline bool Writer<StringBuffer>::ScanWriteUnescapedString(StringStream& is, size_t length) {
if (length < 16)
return RAPIDJSON_LIKELY(is.Tell() < length);
if (!RAPIDJSON_LIKELY(is.Tell() < length))
return false;
const char* p = is.src_;
const char* end = is.head_ + length;
const char* nextAligned = reinterpret_cast<const char*>((reinterpret_cast<size_t>(p) + 15) & static_cast<size_t>(~15));
const char* endAligned = reinterpret_cast<const char*>(reinterpret_cast<size_t>(end) & static_cast<size_t>(~15));
if (nextAligned > end)
return true;
os_->Reserve(nextAligned - p);
while (p != nextAligned)
if (*p < 0x20 || *p == '\"' || *p == '\\') {
is.src_ = p;
return RAPIDJSON_LIKELY(is.Tell() < length);
}
else
os_->PutUnsafe(*p++);
// The rest of string using SIMD
static const char dquote[16] = { '\"', '\"', '\"', '\"', '\"', '\"', '\"', '\"', '\"', '\"', '\"', '\"', '\"', '\"', '\"', '\"' };
static const char bslash[16] = { '\\', '\\', '\\', '\\', '\\', '\\', '\\', '\\', '\\', '\\', '\\', '\\', '\\', '\\', '\\', '\\' };
static const char space[16] = { 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F };
const __m128i dq = _mm_loadu_si128(reinterpret_cast<const __m128i *>(&dquote[0]));
const __m128i bs = _mm_loadu_si128(reinterpret_cast<const __m128i *>(&bslash[0]));
const __m128i sp = _mm_loadu_si128(reinterpret_cast<const __m128i *>(&space[0]));
for (; p != endAligned; p += 16) {
const __m128i s = _mm_load_si128(reinterpret_cast<const __m128i *>(p));
const __m128i t1 = _mm_cmpeq_epi8(s, dq);
const __m128i t2 = _mm_cmpeq_epi8(s, bs);
const __m128i t3 = _mm_cmpeq_epi8(_mm_max_epu8(s, sp), sp); // s < 0x20 <=> max(s, 0x1F) == 0x1F
const __m128i x = _mm_or_si128(_mm_or_si128(t1, t2), t3);
unsigned short r = static_cast<unsigned short>(_mm_movemask_epi8(x));
if (RAPIDJSON_UNLIKELY(r != 0)) { // some of characters is escaped
SizeType len;
#ifdef _MSC_VER // Find the index of first escaped
unsigned long offset;
_BitScanForward(&offset, r);
len = offset;
#else
len = static_cast<SizeType>(__builtin_ffs(r) - 1);
#endif
char* q = reinterpret_cast<char*>(os_->Push(len));
for (size_t i = 0; i < len; i++)
q[i] = p[i];
p += len;
break;
}
_mm_storeu_si128(reinterpret_cast<__m128i *>(os_->Push(16)), s);
}
is.src_ = p;
return RAPIDJSON_LIKELY(is.Tell() < length);
}
#elif defined(RAPIDJSON_NEON) && !defined(__SANITIZE_ADDRESS__)
template<>
inline bool Writer<StringBuffer>::ScanWriteUnescapedString(StringStream& is, size_t length) {
if (length < 16)
return RAPIDJSON_LIKELY(is.Tell() < length);
if (!RAPIDJSON_LIKELY(is.Tell() < length))
return false;
const char* p = is.src_;
const char* end = is.head_ + length;
const char* nextAligned = reinterpret_cast<const char*>((reinterpret_cast<size_t>(p) + 15) & static_cast<size_t>(~15));
const char* endAligned = reinterpret_cast<const char*>(reinterpret_cast<size_t>(end) & static_cast<size_t>(~15));
if (nextAligned > end)
return true;
os_->Reserve(nextAligned - p);
while (p != nextAligned)
if (*p < 0x20 || *p == '\"' || *p == '\\') {
is.src_ = p;
return RAPIDJSON_LIKELY(is.Tell() < length);
}
else
os_->PutUnsafe(*p++);
// The rest of string using SIMD
const uint8x16_t s0 = vmovq_n_u8('"');
const uint8x16_t s1 = vmovq_n_u8('\\');
const uint8x16_t s2 = vmovq_n_u8('\b');
const uint8x16_t s3 = vmovq_n_u8(32);
for (; p != endAligned; p += 16) {
const uint8x16_t s = vld1q_u8(reinterpret_cast<const uint8_t *>(p));
uint8x16_t x = vceqq_u8(s, s0);
x = vorrq_u8(x, vceqq_u8(s, s1));
x = vorrq_u8(x, vceqq_u8(s, s2));
x = vorrq_u8(x, vcltq_u8(s, s3));
x = vrev64q_u8(x); // Rev in 64
uint64_t low = vgetq_lane_u64(vreinterpretq_u64_u8(x), 0); // extract
uint64_t high = vgetq_lane_u64(vreinterpretq_u64_u8(x), 1); // extract
SizeType len = 0;
bool escaped = false;
if (low == 0) {
if (high != 0) {
uint32_t lz = internal::clzll(high);
len = 8 + (lz >> 3);
escaped = true;
}
} else {
uint32_t lz = internal::clzll(low);
len = lz >> 3;
escaped = true;
}
if (RAPIDJSON_UNLIKELY(escaped)) { // some of characters is escaped
char* q = reinterpret_cast<char*>(os_->Push(len));
for (size_t i = 0; i < len; i++)
q[i] = p[i];
p += len;
break;
}
vst1q_u8(reinterpret_cast<uint8_t *>(os_->Push(16)), s);
}
is.src_ = p;
return RAPIDJSON_LIKELY(is.Tell() < length);
}
#endif // RAPIDJSON_NEON
RAPIDJSON_NAMESPACE_END
#if defined(_MSC_VER) || defined(__clang__)
RAPIDJSON_DIAG_POP
#endif
#endif // RAPIDJSON_RAPIDJSON_H_

19
tst/CMakeLists.txt Normal file
View File

@ -0,0 +1,19 @@
##################################################
# We use GoogleTest for both unit and system tests
##################################################
message("tst/CMakeLists.txt")
# Fetch GoogleTest.
include(FetchContent)
FetchContent_Declare(
googletest
GIT_REPOSITORY https://github.com/google/googletest.git
GIT_TAG 58d77fa8070e8cec2dc1ed015d66b454c8d78850 # release-1.12.1
OVERRIDE_FIND_PACKAGE)
FetchContent_MakeAvailable(googletest)
include(GoogleTest)
add_subdirectory(unit)

10
tst/integration/README.md Normal file
View File

@ -0,0 +1,10 @@
# Integration Tests
This directory contains integration tests that verify the interaction between vlkaye-server and valkeyJSON module features working together. Unlike unit tests that test individual components in isolation, these tests validate the system's behavior as a whole.
## Requirements
```text
python 3.9
pytest 4
```

80
tst/integration/data/store.json Normal file
View File

@ -0,0 +1,80 @@
{
"store": {
"books": [
{
"category": "reference",
"author": "Nigel Rees",
"title": "Sayings of the Century",
"price": 8.95,
"in-stock": true
},
{
"category": "fiction",
"author": "Evelyn Waugh",
"title": "Sword of Honour",
"price": 12.99,
"in-stock": true,
"movies": [
{
"title": "Sword of Honour - movie",
"realisator": {
"first_name": "Bill",
"last_name": "Anderson"
}
}
]
},
{
"category": "fiction",
"author": "Herman Melville",
"title": "Moby Dick",
"isbn": "0-553-21311-3",
"price": 9,
"in-stock": false
},
{
"category": "fiction",
"author": "J. R. R. Tolkien",
"title": "The Lord of the Rings",
"isbn": "0-115-03266-2",
"price": 22.99,
"in-stock": true
},
{
"category": "reference",
"author": "William Jr. Strunk",
"title": "The Elements of Style",
"price": 6.99,
"in-stock": false
},
{
"category": "fiction",
"author": "Leo Tolstoy",
"title": "Anna Karenina",
"price": 22.99,
"in-stock": true
},
{
"category": "reference",
"author": "Sarah Janssen",
"title": "The World Almanac and Book of Facts 2021",
"isbn": "0-925-23305-2",
"price": 10.69,
"in-stock": false
},
{
"category": "reference",
"author": "Kate L. Turabian",
"title": "Manual for Writers of Research Papers",
"isbn": "0-675-16695-1",
"price": 8.59,
"in-stock": true
}
],
"bicycle": {
"color": "red",
"price": 19.64,
"in-stock": true
}
}
}

26
tst/integration/data/wikipedia.json Normal file
View File

@ -0,0 +1,26 @@
{
"firstName": "John",
"lastName": "Smith",
"age": 27,
"weight": 135.17,
"isAlive": true,
"address": {
"street": "21 2nd Street",
"city": "New York",
"state": "NY",
"zipcode": "10021-3100"
},
"phoneNumbers": [
{
"type": "home",
"number": "212 555-1234"
},
{
"type": "office",
"number": "646 555-4567"
}
],
"children": [],
"spouse": null,
"groups": {}
}

1
tst/integration/data/wikipedia_compact.json Normal file
View File

@ -0,0 +1 @@
{"firstName":"John","lastName":"Smith","age":27,"weight":135.17,"isAlive":true,"address":{"street":"21 2nd Street","city":"New York","state":"NY","zipcode":"10021-3100"},"phoneNumbers":[{"type":"home","number":"212 555-1234"},{"type":"office","number":"646 555-4567"}],"children":[],"spouse":null,"groups":{}}

31
tst/integration/error_handlers.py Normal file
View File

@ -0,0 +1,31 @@
import re
class ErrorStringTester:
def is_syntax_error(string):
return string.startswith("SYNTAXERR") or \
string.startswith("unknown subcommand")
def is_nonexistent_error(string):
return string.startswith("NONEXISTENT")
def is_wrongtype_error(string):
return string.startswith("WRONGTYPE")
def is_number_overflow_error(string):
return string.startswith("OVERFLOW")
def is_outofboundaries_error(string):
return string.startswith("OUTOFBOUNDARIES")
def is_limit_exceeded_error(string):
return string.startswith("LIMIT")
def is_write_error(string):
return string.startswith("ERROR") or string.startswith("OUTOFBOUNDARIES") or \
string.startswith("WRONGTYPE") or string.startswith("NONEXISTENT")
# NOTE: Uses .find instead of .startswith in case prefix added in the future
def is_wrong_number_of_arguments_error(string):
return string.find("wrong number of arguments") >= 0 or \
string.lower().find('invalid number of arguments') >= 0

61
tst/integration/json_test_case.py Normal file
View File

@ -0,0 +1,61 @@
import valkey
import pytest
import os
import logging
import shutil
import time
from valkeytests.valkey_test_case import ValkeyTestCase, ValkeyServerHandle
from valkey import ResponseError
from error_handlers import ErrorStringTester
class SimpleTestCase(ValkeyTestCase):
'''
Simple test case, single server without loading JSON module.
'''
def setup(self):
super(SimpleTestCase, self).setup()
self.client = self.server.get_new_client()
def teardown(self):
if self.is_connected():
self.client.execute_command("FLUSHALL")
logging.info("executed FLUSHALL at teardown")
super(SimpleTestCase, self).teardown()
def is_connected(self):
try:
self.client.ping()
return True
except:
return False
class JsonTestCase(SimpleTestCase):
'''
Base class for JSON test, single server with JSON module loaded.
'''
def get_custom_args(self):
self.set_server_version(os.environ['SERVER_VERSION'])
return {
'loadmodule': os.getenv('MODULE_PATH'),
"enable-debug-command": "local",
'enable-protected-configs': 'yes'
}
def verify_error_response(self, client, cmd, expected_err_reply):
try:
client.execute_command(cmd)
assert False
except ResponseError as e:
assert_error_msg = f"Actual error message: '{str(e)}' is different from expected error message '{expected_err_reply}'"
assert str(e) == expected_err_reply, assert_error_msg
def setup(self):
super(JsonTestCase, self).setup()
self.error_class = ErrorStringTester
def teardown(self):
super(JsonTestCase, self).teardown()

34
tst/integration/run.sh Executable file
View File

@ -0,0 +1,34 @@
#!/usr/bin/env bash
# Sometimes processes are left running when test is cancelled.
# Therefore, before build start, we kill all running test processes left from previous test run.
echo "Kill old running test"
pkill -9 -x Pytest || true
pkill -9 -f "valkey-server.*:" || true
pkill -9 -f Valgrind || true
pkill -9 -f "valkey-benchmark" || true
# cd to the current directory of the script
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
cd "${DIR}"
export MODULE_PATH=$2/build/src/libjson.so
echo "Running integration tests against Valkey version: $SERVER_VERSION"
if [[ ! -z "${TEST_PATTERN}" ]] ; then
export TEST_PATTERN="-k ${TEST_PATTERN}"
fi
BINARY_PATH=".build/binaries/$SERVER_VERSION/valkey-server"
if [[ ! -f "${BINARY_PATH}" ]] ; then
echo "${BINARY_PATH} missing"
exit 1
fi
if [[ $1 == "test" ]] ; then
python -m pytest --html=report.html --cache-clear -v ${TEST_FLAG} ./ ${TEST_PATTERN}
else
echo "Unknown target: $1"
exit 1
fi

4104
tst/integration/test_json_basic.py Normal file
View File

File diff suppressed because it is too large Load Diff

44
tst/integration/test_rdb.py Normal file
View File

@ -0,0 +1,44 @@
from utils_json import DEFAULT_MAX_PATH_LIMIT, \
DEFAULT_STORE_PATH
from valkey.exceptions import ResponseError, NoPermissionError
from valkeytests.conftest import resource_port_tracker
import pytest
import glob
import logging
import os
import random
import struct
import json
from math import isclose, isnan, isinf, frexp
from json_test_case import JsonTestCase
class TestRdb(JsonTestCase):
def setup_data(self):
client = self.server.get_new_client()
client.config_set(
'json.max-path-limit', DEFAULT_MAX_PATH_LIMIT)
# Need the following line when executing the test against a running Valkey.
# Otherwise, data from previous test cases will interfere current test case.
client.execute_command("FLUSHDB")
# Load strore sample JSONs. We use strore.json as input to create a document key. Then, use
# strore_compact.json, which does not have indent/space/newline, to verify correctness of serialization.
with open(DEFAULT_STORE_PATH, 'r') as file:
self.data_store = file.read()
assert b'OK' == client.execute_command(
'JSON.SET', 'store', '.', self.data_store)
def setup(self):
super(TestRdb, self).setup()
self.setup_data()
def test_rdb_saverestore(self):
"""
Test RDB saving
"""
client = self.server.get_new_client()
assert True == client.execute_command('save')
client.execute_command('FLUSHDB')
assert b'OK' == client.execute_command('DEBUG', 'RELOAD', 'NOSAVE')

28
tst/integration/utils_json.py Normal file
View File

@ -0,0 +1,28 @@
import pytest
import os
import random
import string
from valkey.exceptions import ResponseError
JSON_MODULE_NAME = 'json'
JSON_INFO_NAMES = {
'num_documents': JSON_MODULE_NAME + '_num_documents',
'total_memory_bytes': JSON_MODULE_NAME + '_total_memory_bytes',
'doc_histogram': JSON_MODULE_NAME + '_doc_histogram',
'read_histogram': JSON_MODULE_NAME + '_read_histogram',
'insert_histogram': JSON_MODULE_NAME + '_insert_histogram',
'update_histogram': JSON_MODULE_NAME + '_update_histogram',
'delete_histogram': JSON_MODULE_NAME + '_delete_histogram',
'max_path_depth_ever_seen': JSON_MODULE_NAME + '_max_path_depth_ever_seen',
'max_document_size_ever_seen': JSON_MODULE_NAME + '_max_document_size_ever_seen',
'total_malloc_bytes_used': JSON_MODULE_NAME + "_total_malloc_bytes_used",
'memory_traps_enabled': JSON_MODULE_NAME + "_memory_traps_enabled",
}
DEFAULT_MAX_DOCUMENT_SIZE = 64*1024*1024
DEFAULT_MAX_PATH_LIMIT = 128
DEFAULT_WIKIPEDIA_PATH = 'data/wikipedia.json'
DEFAULT_WIKIPEDIA_COMPACT_PATH = 'data/wikipedia_compact.json'
DEFAULT_STORE_PATH = 'data/store.json'
JSON_INFO_METRICS_SECTION = JSON_MODULE_NAME + '_core_metrics'
JSON_MODULE_NAME = 'json'

89
tst/unit/CMakeLists.txt Normal file
View File

@ -0,0 +1,89 @@
#########################################
# Define unit tests
#########################################
message("tst/unit/CMakeLists.txt: Define unit tests")
# This is the set of sources for the basic test
file(GLOB_RECURSE UNIT_TEST_SRC "*.cc")
#########################################
# Tell Cmake how to run the unit tests
#########################################
# A brief philosophical thought, about unit tests: if possible, it's preferable to have all unit
# tests in a single (or a low number of) binary executable. This is disk-space efficient for the
# test suite, avoids unnecessary linking steps, and provides a nice, simple way to interface with
# the test suite (should you need to do so manually, or, for instance, with a debugger). Large
# numbers of test binaries are certainly possible, and in some rare cases are even necessary, but
# don't provide many advantages over a single binary in the average case.
# This file defines a single test executable, "unitTests", which uses the Googletest framework.
# This defines each unit test and associates it with its sources
add_executable(unitTests ${UNIT_TEST_SRC})
# Build with C11 & C++17
set_target_properties(
unitTests
PROPERTIES
C_STANDARD 11
C_STANDARD_REQUIRED ON
CXX_STANDARD 17
CXX_STANDARD_REQUIRED ON
POSITION_INDEPENDENT_CODE ON
)
target_include_directories(unitTests
PRIVATE
${PROJECT_SOURCE_DIR}/src
${rapidjson_SOURCE_DIR}/include
)
# Add dependency to the code under test.
target_link_libraries(unitTests ${JSON_MODULE_LIB})
# Link GoogleTest libraries after fetch
target_link_libraries(unitTests
GTest::gtest_main # Link the main GoogleTest library
GTest::gmock_main # Link GoogleMock
)
# This tells CTest about this unit test executable
# The TEST_PREFIX prepends "unit_" to the name of these tests in the output,
# which makes them easier to identify at a glance
# The TEST_LIST settings creates a CMake list of all of the tests in the
# binary. This is useful for, for instance, the set_tests_properties statement
# below
# For more information, see: https://cmake.org/cmake/help/v3.12/module/GoogleTest.html
# To get this to work properly in a cross-compile environment, you need to set up
# CROSSCOMPILING_EMULATOR (see https://cmake.org/cmake/help/v3.12/prop_tgt/CROSSCOMPILING_EMULATOR.html)
# DISCOVERY_TIMEOUT - number of seconds given for Gtest to discover the tests to run, it should
# be big enough so the tests can start on MacOS and can be any number, 59 is just prime number
# close to 1 minute ;)
gtest_discover_tests(unitTests
TEST_PREFIX unit_
TEST_LIST unit_gtests
DISCOVERY_TIMEOUT 59
)
# This tells the CTest harness about how it should treat these tests. For
# instance, you can uncomment the RUN_SERIAL line to force the tests to run
# sequentially (e.g. if the tests are not thread-safe... in most cases, tests
# SHOULD be thread-safe). We also set a high-level timeout: if the test takes
# longer than the specified time, it is killed by the harness and reported as a
# failure. And finally, we provide a "label" that is used by CTest when
# reporting result statistics (e.g. "UnitTests: 72 successes, 3 failures").
# For more properties that can be set, see:
# https://cmake.org/cmake/help/v3.9/manual/cmake-properties.7.html#test-properties
set_tests_properties(${unit_gtests} PROPERTIES
# RUN_SERIAL 1
TIMEOUT 10 # seconds
LABELS UnitTests
)
add_custom_target(unit
COMMAND ${CMAKE_BINARY_DIR}/tst/unit/unitTests
DEPENDS unitTests
WORKING_DIRECTORY ${CMAKE_BINARY_DIR}
COMMENT "Running unit tests..."
)

6
tst/unit/CPPLINT.cfg Normal file
View File

@ -0,0 +1,6 @@
filter=-build/include_subdir
# STL allocator needs implicit single arg constructor which CPPLINT doesn't like by default
filter=-runtime/explicit
filter=-runtime/string
filter=-runtime/int
linelength=120

2304
tst/unit/dom_test.cc Normal file
View File

File diff suppressed because it is too large Load Diff

209
tst/unit/hashtable_test.cc Normal file
View File

@ -0,0 +1,209 @@
#include <cstdlib>
#include <cstddef>
#include <cstring>
#include <cstdio>
#include <cstdint>
#include <memory>
#include <deque>
#include <map>
#include <string>
#include <sstream>
#include <random>
#include <limits>
#include <vector>
#include <cmath>
#include <utility>
#include <gtest/gtest.h>
#include "json/dom.h"
#include "json/alloc.h"
#include "json/stats.h"
#include "json/keytable.h"
#include "module_sim.h"
// Cheap, predictable hash
static size_t hash1(const char *ptr, size_t len) {
(void)ptr;
return len;
}
class HashTableTest : public ::testing::Test {
protected:
void SetUp() override {
}
size_t original_malloced;
void TearDown() override {
if (keyTable) {
malloced = original_malloced;
EXPECT_EQ(keyTable->validate(), "");
}
delete keyTable;
}
void Setup1(size_t numShards = 1, size_t htsize = 0, size_t (*h)(const char *, size_t) = hash1) {
setupValkeyModulePointers();
KeyTable::Config c;
c.malloc = dom_alloc;
c.free = dom_free;
c.hash = h;
c.numShards = numShards;
keyTable = new KeyTable(c);
rapidjson::hashTableFactors.minHTSize = htsize;
original_malloced = malloced;
malloced = 0; // Ignore startup memory consumption
}
};
TEST_F(HashTableTest, simple) {
Setup1();
{
JValue v;
v.SetObject();
v.AddMember(JValue("True"), JValue(true), allocator);
EXPECT_EQ(v.MemberCount(), 1u);
EXPECT_TRUE(v["True"].IsBool());
EXPECT_GT(malloced, 0);
}
EXPECT_EQ(malloced, 0);
}
static JValue makeKey(size_t i) {
return std::move(JValue().SetString(std::to_string(i), allocator));
}
static JValue makeArray(size_t sz, size_t offset = 0) {
JValue j;
j.SetArray();
for (size_t i = 0; i < sz; ++i) {
j.PushBack(JValue(i + offset), allocator);
}
return j;
}
static JValue makeArrayArray(size_t p, size_t q) {
JValue j = makeArray(p);
for (size_t i = 0; i < p; ++i) {
j[i] = makeArray(q, i);
}
return j;
}
TEST_F(HashTableTest, checkeq) {
Setup1();
for (size_t i : {0, 1, 10}) {
ASSERT_EQ(makeArrayArray(i, i), makeArrayArray(i, i));
}
}
TEST_F(HashTableTest, insertAndRemoveMany) {
Setup1(1, 5);
for (size_t sz : {10, 50, 100}) {
EXPECT_EQ(malloced, 0);
rapidjson::hashTableStats.reset();
{
JValue v;
v.SetObject();
EXPECT_EQ(v.Validate(), "");
for (size_t i = 0; i < sz; ++i) {
v.AddMember(makeKey(i), makeArrayArray(i, i), allocator);
EXPECT_EQ(v.Validate(), "");
}
EXPECT_EQ(v.MemberCount(), sz);
EXPECT_GT(rapidjson::hashTableStats.rehashUp, 0);
EXPECT_EQ(rapidjson::hashTableStats.convertToHT, 1);
auto s = keyTable->getStats();
EXPECT_EQ(s.size, sz);
for (size_t i = 0; i < sz; ++i) EXPECT_EQ(v[makeKey(i)], makeArrayArray(i, i));
for (size_t i = 0; i < sz; ++i) {
v.RemoveMember(makeKey(i));
EXPECT_EQ(v.Validate(), "");
}
EXPECT_GT(rapidjson::hashTableStats.rehashDown, 0);
EXPECT_EQ(v.MemberCount(), 0);
s = keyTable->getStats();
EXPECT_EQ(s.size, 0); // All entries should be gone.
}
EXPECT_EQ(malloced, 0);
}
}
TEST_F(HashTableTest, SetObjectRawHT) {
Setup1();
std::ostringstream os;
os << "{\"a\":1";
for (size_t i = 0; i < 100; ++i) os << ",\"" << i << "\":" << i;
os << "}";
JDocument *doc;
JsonUtilCode rc = dom_parse(nullptr, os.str().c_str(), os.str().size(), &doc);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(rapidjson::hashTableStats.reserveHT, 1);
rapidjson::StringBuffer oss;
dom_serialize(doc, nullptr, oss);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(oss.GetString(), os.str());
dom_free_doc(doc);
auto s = keyTable->getStats();
EXPECT_EQ(s.size, 0); // All entries should be gone.
EXPECT_EQ(malloced, 0);
}
TEST_F(HashTableTest, CopyMembers) {
Setup1(1, 5);
for (size_t sz : {10, 50, 100}) {
rapidjson::hashTableStats.reset();
JValue v;
v.SetObject();
EXPECT_EQ(v.Validate(), "");
for (size_t i = 0; i < sz; ++i) {
v.AddMember(makeKey(i), makeArrayArray(i, i), allocator);
EXPECT_EQ(v.Validate(), "");
}
EXPECT_EQ(v.MemberCount(), sz);
EXPECT_GT(rapidjson::hashTableStats.rehashUp, 0);
EXPECT_EQ(rapidjson::hashTableStats.convertToHT, 1);
auto s = keyTable->getStats();
EXPECT_EQ(s.size, sz);
EXPECT_EQ(s.handles, sz);
{
rapidjson::hashTableStats.reset();
JValue v2(v, allocator); // Invokes copymembers
EXPECT_EQ(v2.Validate(), "");
EXPECT_EQ(v2.MemberCount(), sz);
EXPECT_EQ(rapidjson::hashTableStats.rehashUp, 0);
EXPECT_EQ(rapidjson::hashTableStats.rehashDown, 0);
EXPECT_EQ(rapidjson::hashTableStats.convertToHT, 0);
s = keyTable->getStats();
EXPECT_EQ(s.size, sz);
EXPECT_EQ(s.handles, sz*2);
for (size_t i = 0; i < sz; ++i) {
EXPECT_EQ(v[makeKey(i)].GetArray(), makeArrayArray(i, i));
EXPECT_EQ(v2[makeKey(i)].GetArray(), makeArrayArray(i, i));
}
}
}
EXPECT_EQ(malloced, 0);
}
//
// Test that hash tables > 2^19 are properly handled.
//
TEST_F(HashTableTest, DistributionTest) {
extern size_t hash_function(const char *, size_t);
Setup1(1, 0, hash_function);
enum { TABLE_SIZE_BITS = 22 }; // LOG2(Table Size)
enum { TABLE_SIZE = 1ull << TABLE_SIZE_BITS };
JValue v;
v.SetObject();
for (size_t i = 0; i < TABLE_SIZE; ++i) {
v.AddMember(makeKey(i), JValue(true), allocator);
}
//
// Now, compute the distribution stats, make sure the longest run is sufficiently small
//
std::map<size_t, size_t> runs;
v.getObjectDistribution(runs, 5);
// std::cout << "Dist:";
// for (auto& x : runs) std::cout << x.first << ":" << x.second << ",";
// std::cout << std::endl;
ASSERT_NE(runs.size(), 0u);
EXPECT_LT(runs.rbegin()->first, 0.0001 * TABLE_SIZE);
}

249
tst/unit/json_test.cc Normal file
View File

@ -0,0 +1,249 @@
#include <cstdlib>
#include <cstddef>
#include <cstring>
#include <cstdio>
#include <cstdint>
#include <cmath>
#include <memory>
#include <deque>
#include <vector>
#include <string>
#include <sstream>
#include <utility>
#include <iostream>
#include <unordered_map>
#include <map>
#include <set>
#include <gtest/gtest.h>
#include "json/dom.h"
#include "json/alloc.h"
#include "json/stats.h"
#include "json/selector.h"
extern void SetupAllocFuncs(size_t numShards);
class JsonTest : public ::testing::Test {
void SetUp() override {
JsonUtilCode rc = jsonstats_init();
ASSERT_EQ(rc, JSONUTIL_SUCCESS);
SetupAllocFuncs(16);
}
void TearDown() override {
delete keyTable;
keyTable = nullptr;
}
void setShards(size_t numShards) {
if (keyTable) delete keyTable;
SetupAllocFuncs(numShards);
}
};
TEST_F(JsonTest, testArrIndex_fullobjects) {
const char *input = "[5, 6, {\"a\":\"b\"}, [99,100], [\"c\"]]";
JDocument *doc;
JsonUtilCode rc = dom_parse(nullptr, input, strlen(input), &doc);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
jsn::vector<int64_t> indexes;
bool is_v2_path;
rc = dom_array_index_of(doc, ".", "{\"a\":\"b\"}", 9, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_FALSE(is_v2_path);
EXPECT_EQ(indexes.size(), 1);
EXPECT_EQ(indexes[0], 2);
rc = dom_array_index_of(doc, ".", "[\"c\"]", 5, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_FALSE(is_v2_path);
EXPECT_EQ(indexes.size(), 1);
EXPECT_EQ(indexes[0], 4);
rc = dom_array_index_of(doc, ".", "[99,100]", 8, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_FALSE(is_v2_path);
EXPECT_EQ(indexes.size(), 1);
EXPECT_EQ(indexes[0], 3);
dom_free_doc(doc);
}
TEST_F(JsonTest, testArrIndex_arr) {
const char *input = "{\"a\":[1,2,[15,50],3], \"nested\": {\"a\": [3,4,[5,5]]}}";
JDocument *doc;
JsonUtilCode rc = dom_parse(nullptr, input, strlen(input), &doc);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
jsn::vector<int64_t> indexes;
bool is_v2_path;
rc = dom_array_index_of(doc, "$..a", "[15,50]", 7, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 2);
EXPECT_EQ(indexes[0], 2);
EXPECT_EQ(indexes[1], -1);
rc = dom_array_index_of(doc, "$..a", "3", 1, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 2);
EXPECT_EQ(indexes[0], 3);
EXPECT_EQ(indexes[1], 0);
rc = dom_array_index_of(doc, "$..a", "[5,5]", 5, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 2);
EXPECT_EQ(indexes[0], -1);
EXPECT_EQ(indexes[1], 2);
rc = dom_array_index_of(doc, "$..a", "35", 2, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 2);
EXPECT_EQ(indexes[0], -1);
EXPECT_EQ(indexes[0], -1);
dom_free_doc(doc);
}
TEST_F(JsonTest, testArrIndex_object) {
const char *input = "{\"a\":{\"b\":[2,4,{\"a\":4},false,true,{\"b\":false}]}}";
JDocument *doc;
JsonUtilCode rc = dom_parse(nullptr, input, strlen(input), &doc);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
jsn::vector<int64_t> indexes;
bool is_v2_path;
rc = dom_array_index_of(doc, "$.a.b", "{\"a\":4}", 7, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 1);
EXPECT_EQ(indexes[0], 2);
rc = dom_array_index_of(doc, "$.a.b", "{\"b\":false}", 11, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 1);
EXPECT_EQ(indexes[0], 5);
rc = dom_array_index_of(doc, "$.a.b", "false", 5, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 1);
EXPECT_EQ(indexes[0], 3);
rc = dom_array_index_of(doc, "$..a", "{\"a\":4}", 7, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 2);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], INT64_MAX);
rc = dom_array_index_of(doc, "$..a..", "{\"a\":4}", 7, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 4);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], 2);
EXPECT_EQ(indexes[2], INT64_MAX);
EXPECT_EQ(indexes[3], INT64_MAX);
dom_free_doc(doc);
}
TEST_F(JsonTest, testArrIndex_nested_search) {
const char *input = "{\"level0\":{\"level1_0\":{\"level2\":"
"[1,2,3, [25, [4,5,{\"c\":\"d\"}]]]},"
"\"level1_1\":{\"level2\": [[{\"a\":[2,5]}, true, null]]}}}";
JDocument *doc;
JsonUtilCode rc = dom_parse(nullptr, input, strlen(input), &doc);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
jsn::vector<int64_t> indexes;
bool is_v2_path;
rc = dom_array_index_of(doc, "$..level0.level1_0..", "[4,5,{\"c\":\"d\"}]", 15, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 5);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], -1);
EXPECT_EQ(indexes[2], 1);
EXPECT_EQ(indexes[3], -1);
EXPECT_EQ(indexes[4], INT64_MAX);
rc = dom_array_index_of(doc, "$..level0.level1_0..", "[25, [4,5,{\"c\":\"d\"}]]", 21, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 5);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], 3);
EXPECT_EQ(indexes[2], -1);
EXPECT_EQ(indexes[3], -1);
EXPECT_EQ(indexes[4], INT64_MAX);
rc = dom_array_index_of(doc, "$..level0.level1_0..", "{\"c\":\"d\"}", 9, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 5);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], -1);
EXPECT_EQ(indexes[2], -1);
EXPECT_EQ(indexes[3], 2);
EXPECT_EQ(indexes[4], INT64_MAX);
rc = dom_array_index_of(doc, "$..level0.level1_0..", "[4,5,{\"a\":\"b\"}]", 15, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 5);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], -1);
EXPECT_EQ(indexes[2], -1);
EXPECT_EQ(indexes[3], -1);
EXPECT_EQ(indexes[4], INT64_MAX);
rc = dom_array_index_of(doc, "$..level0.level1_1..", "[null,true,{\"a\":[2,5]}]", 23, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 5);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], -1);
EXPECT_EQ(indexes[2], -1);
EXPECT_EQ(indexes[3], INT64_MAX);
EXPECT_EQ(indexes[4], -1);
rc = dom_array_index_of(doc, "$..level0.level1_1..", "[{\"a\":[2,5]},true,null]", 23, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 5);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], 0);
EXPECT_EQ(indexes[2], -1);
EXPECT_EQ(indexes[3], INT64_MAX);
EXPECT_EQ(indexes[4], -1);
rc = dom_array_index_of(doc, "$..level0.level1_1..", "[{\"a\":[2,5]},true]", 18, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 5);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], -1);
EXPECT_EQ(indexes[2], -1);
EXPECT_EQ(indexes[3], INT64_MAX);
EXPECT_EQ(indexes[4], -1);
rc = dom_array_index_of(doc, "$..level0.level1_0..", "[4,{\"c\":\"d\"}]", 13, 0, 0, indexes, is_v2_path);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_TRUE(is_v2_path);
EXPECT_EQ(indexes.size(), 5);
EXPECT_EQ(indexes[0], INT64_MAX);
EXPECT_EQ(indexes[1], -1);
EXPECT_EQ(indexes[2], -1);
EXPECT_EQ(indexes[3], -1);
EXPECT_EQ(indexes[4], INT64_MAX);
dom_free_doc(doc);
}

393
tst/unit/keytable_test.cc Normal file
View File

@ -0,0 +1,393 @@
#include <cstdlib>
#include <cstddef>
#include <cstring>
#include <cstdio>
#include <cstdint>
#include <memory>
#include <deque>
#include <string>
#include <sstream>
#include <random>
#include <limits>
#include <vector>
#include <cmath>
#include <gtest/gtest.h>
#include "json/dom.h"
#include "json/alloc.h"
#include "json/stats.h"
#include "json/keytable.h"
#include "json/memory.h"
#include "module_sim.h"
class PtrWithMetaDataTest : public ::testing::Test {
};
TEST_F(PtrWithMetaDataTest, t) {
memory_traps_control(false); // Necessary so that MEMORY_VALIDATE buried in getPointer doesn't croak on bad memory
EXPECT_EQ(0x7FFFF, PtrWithMetaData<size_t>::METADATA_MASK);
for (size_t i = 1; i & 0x7FFFF; i <<= 1) {
size_t var = 0xdeadbeeffeedfeddull;
PtrWithMetaData<size_t> p(&var, i);
EXPECT_EQ(&*p, &var);
EXPECT_EQ(*p, var);
EXPECT_EQ(size_t(p.getMetaData()), i);
p.clear();
EXPECT_EQ(p.getMetaData(), 0);
p.setMetaData(i);
EXPECT_EQ(size_t(p.getMetaData()), i);
}
for (size_t i = 8; i & 0x0000FFFFFFFFFFF8ull; i <<= 1) {
PtrWithMetaData<size_t> p(reinterpret_cast<size_t *>(i), 0x7FFFF);
EXPECT_EQ(size_t(&*p), i);
EXPECT_EQ(p.getMetaData(), 0x7FFFF);
}
}
// Cheap, predictable hash
static size_t hash1(const char *ptr, size_t len) {
(void)ptr;
return len;
}
extern size_t MAX_FAST_TABLE_SIZE; // in keytable.cc
class KeyTableTest : public ::testing::Test {
protected:
void SetUp() override {
}
void TearDown() override {
if (t) {
EXPECT_EQ(t->validate(), "");
}
delete t;
}
void Setup1(size_t numShards = 1, size_t (*hf)(const char *, size_t) = hash1) {
setupValkeyModulePointers();
KeyTable::Config c;
c.malloc = dom_alloc;
c.free = dom_free;
c.hash = hf;
c.numShards = numShards;
t = new KeyTable(c);
}
KeyTable *t = nullptr;
};
TEST_F(KeyTableTest, layoutTest) {
Setup1();
size_t bias = 10;
for (size_t slen : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
0xFF, 0x100, 0xFFFF, 0x10000, 0xFFFFFF, 0x1000000}) {
std::string s;
s.resize(slen);
for (size_t i = 0; i < slen; ++i) {
s[i] = i + bias;
}
KeyTable_Layout *l = KeyTable_Layout::makeLayout(dom_alloc, s.data(), s.length(), 0, false);
ASSERT_EQ(l->getLength(), slen);
for (size_t i = 0; i < slen; ++i) {
ASSERT_EQ(0xFF & (i + bias), 0xFF & (l->getText()[i]));
}
dom_free(l);
bias++;
}
}
TEST_F(KeyTableTest, testInitialization) {
Setup1();
EXPECT_EQ(t->validate(), "");
EXPECT_GT(malloced, 0);
EXPECT_EQ(t->validate(), "");
auto s = t->getStats();
EXPECT_EQ(s.size, 0);
EXPECT_EQ(s.handles, 0);
EXPECT_EQ(s.bytes, 0);
EXPECT_GT(s.maxTableSize, 0);
EXPECT_GT(s.totalTable, 0);
EXPECT_EQ(s.rehashes, 0);
EXPECT_EQ(s.maxSearch, 0);
delete t;
t = nullptr;
EXPECT_EQ(malloced, 0);
}
TEST_F(KeyTableTest, testDuplication) {
std::string e = "Empty";
Setup1();
auto f = t->getFactors();
f.maxLoad = 1.0; // No rehashes until we're full....
f.minLoad = std::numeric_limits<float>::min();
t->setFactors(f);
KeyTable_Handle h1 = t->makeHandle(e);
EXPECT_TRUE(h1);
EXPECT_EQ(t->validate(), "");
KeyTable_Handle h2 = t->makeHandle(e);
EXPECT_EQ(t->validate(), "");
EXPECT_TRUE(h2);
EXPECT_EQ(h1, h2);
EXPECT_EQ(&*h1, &*h2);
auto s = t->getStats();
EXPECT_EQ(s.size, 1);
EXPECT_EQ(s.handles, 2);
EXPECT_EQ(s.bytes, 5);
t->destroyHandle(h1);
EXPECT_TRUE(!h1);
EXPECT_EQ(t->validate(), "");
s = t->getStats();
EXPECT_EQ(s.size, 1);
EXPECT_EQ(s.handles, 1);
EXPECT_EQ(s.bytes, 5);
t->destroyHandle(h2);
EXPECT_TRUE(!h2);
EXPECT_EQ(t->validate(), "");
s = t->getStats();
EXPECT_EQ(s.rehashes, 0);
}
TEST_F(KeyTableTest, testClone) {
std::string e = "Empty";
Setup1();
auto f = t->getFactors();
f.maxLoad = 1.0; // No rehashes until we're full....
f.minLoad = std::numeric_limits<float>::min();
t->setFactors(f);
KeyTable_Handle h1 = t->makeHandle(e);
EXPECT_TRUE(h1);
EXPECT_EQ(t->validate(), "");
KeyTable_Handle h2 = t->clone(h1);
EXPECT_EQ(t->validate(), "");
EXPECT_TRUE(h2);
EXPECT_EQ(h1, h2);
EXPECT_EQ(&*h1, &*h2);
auto s = t->getStats();
EXPECT_EQ(s.size, 1);
EXPECT_EQ(s.handles, 2);
EXPECT_EQ(s.bytes, 5);
t->destroyHandle(h1);
EXPECT_TRUE(!h1);
EXPECT_EQ(t->validate(), "");
s = t->getStats();
EXPECT_EQ(s.size, 1);
EXPECT_EQ(s.handles, 1);
EXPECT_EQ(s.bytes, 5);
t->destroyHandle(h2);
EXPECT_TRUE(!h2);
EXPECT_EQ(t->validate(), "");
s = t->getStats();
EXPECT_EQ(s.rehashes, 0);
}
TEST_F(KeyTableTest, SimpleRehash) {
Setup1(1); // 4 element table is the minimum.
auto f = t->getFactors();
f.maxLoad = 1.0; // No rehashes until we're full....
f.minLoad = std::numeric_limits<float>::min();
f.grow = 1.0;
f.shrink = 0.5;
t->setFactors(f);
std::vector<KeyTable_Handle> h;
std::vector<std::string> keys;
std::string k = "";
for (size_t i = 0; i < 4; ++i) {
h.push_back(t->makeHandle(k));
keys.push_back(k);
auto s = t->getStats();
EXPECT_EQ(s.size, i+1);
EXPECT_EQ(s.rehashes, 0);
EXPECT_EQ(t->validate(), "");
k += '*';
}
for (size_t i = 4; i < 8; ++i) {
auto f = t->getFactors();
f.maxLoad = i == 4 ? .5 : 1.0; // No rehashes until we're full....
t->setFactors(f);
h.push_back(t->makeHandle(k));
keys.push_back(k);
auto s = t->getStats();
EXPECT_EQ(s.size, i+1);
EXPECT_EQ(s.rehashes, i == 4 ? 1 : 0);
EXPECT_EQ(t->validate(), "");
k += '*';
}
//
// Now shrink
//
for (size_t i = 0; i < 4; ++i) {
t->destroyHandle(h.back());
h.pop_back();
auto s = t->getStats();
EXPECT_EQ(s.rehashes, 0);
EXPECT_EQ(t->validate(), "");
}
// Next destroyHandle should case a rehash.
for (size_t i = 0; i < 4; ++i) {
auto f = t->getFactors();
f.minLoad = i == 0 ? .5f : std::numeric_limits<float>::min(); // No rehashes until we're full....
t->setFactors(f);
t->destroyHandle(h.back());
h.pop_back();
auto s = t->getStats();
EXPECT_EQ(s.maxTableSize, 4);
EXPECT_EQ(s.rehashes, i == 0 ? 1 : 0);
EXPECT_EQ(t->validate(), "");
}
}
//
// Generate some strings, duplicates are ok.
// Because the hash is the length + the last character the total number of unique strings
// is only 10x of the max length (from the random distribution)
//
std::default_random_engine generator(0);
std::uniform_int_distribution dice(0, 10000); // there are actually ~10x this number of unique strings
size_t make_rand() {
return dice(generator);
}
std::string make_key() {
size_t len = make_rand();
size_t lastDigit = make_rand() % 10;
std::string k;
for (size_t i = 0; i < len; ++i) k += '*';
k += '0' + lastDigit;
return k;
}
TEST_F(KeyTableTest, BigTest) {
//
// Make a zillion keys, Yes, there will be lots of duplicates -> Intentionally
//
for (size_t ft : { 1 << 8, 1 << 10, 1 << 12}) {
MAX_FAST_TABLE_SIZE = ft;
for (size_t numShards : {1, 2}) {
for (size_t numKeys : {1000}) {
Setup1(numShards);
auto f = t->getFactors();
f.grow = 1.1; // Grow slowly
f.maxLoad = .95; // Let the table get REALLY full between hashes
t->setFactors(f);
std::vector<KeyTable_Handle> h;
std::vector<std::string> k;
for (size_t i = 0; i < numKeys; ++i) {
k.push_back(make_key());
h.push_back(t->makeHandle(k.back().c_str(), k.back().length()));
if (0 == (i & 0xFF)) {
EXPECT_EQ(t->validate(), "");
}
}
auto s = t->getStats();
EXPECT_EQ(s.handles, k.size());
EXPECT_LT(s.size, k.size()); // must have at least one duplicate
EXPECT_GT(s.rehashes, 5); // should have had several rehashes
//
// now delete them SLOWLY with lots of rehashes
//
f = t->getFactors();
f.shrink = .05; // Shrink slowly
f.minLoad = .9; // Let the table get REALLY full between hashes
t->setFactors(f);
for (size_t i = 0; i < numKeys; ++i) {
t->destroyHandle(h[i]);
if (0 == (i & 0xFF)) {
EXPECT_EQ(t->validate(), "");
}
}
//
// Teardown.
//
EXPECT_EQ(t->validate(), "");
s = t->getStats();
EXPECT_GT(s.rehashes, 10);
EXPECT_EQ(s.size, 0);
delete t;
t = nullptr;
}
}
}
}
TEST_F(KeyTableTest, StuckKeys) {
Setup1(1);
KeyTable_Layout::setMaxRefCount(3);
std::string e = "Empty";
KeyTable_Handle h1 = t->makeHandle(e);
KeyTable_Handle h2 = t->makeHandle(e);
KeyTable_Handle h3 = t->makeHandle(e);
KeyTable_Handle h4 = t->makeHandle(e);
EXPECT_EQ(t->validate(), "");
auto s = t->getStats();
EXPECT_EQ(s.size, 1);
EXPECT_EQ(s.stuckKeys, 1);
EXPECT_EQ(s.handles, 4);
t->destroyHandle(h1);
t->destroyHandle(h2);
t->destroyHandle(h3);
t->destroyHandle(h4);
s = t->getStats();
EXPECT_EQ(s.stuckKeys, 1);
EXPECT_EQ(s.size, 1);
EXPECT_EQ(s.handles, 0);
}
//
// Make a very large shard, check some stats, delete the elements and see if it shrinks
//
extern size_t hash_function(const char *, size_t);
TEST_F(KeyTableTest, BigShard) {
memory_traps_control(false);
Setup1(1, hash_function);
enum { TABLE_SIZE_BITS = 22 }; // LOG2(Table Size)
enum { TABLE_SIZE = 1ull << TABLE_SIZE_BITS };
std::vector<KeyTable_Handle> handles1;
std::vector<KeyTable_Handle> handles2;
//
// Fill up the table
//
for (size_t i = 0; i < TABLE_SIZE; ++i) {
handles1.push_back(t->makeHandle(std::to_string(i)));
}
auto s = t->getStats();
EXPECT_EQ(s.size, TABLE_SIZE);
EXPECT_EQ(s.handles, TABLE_SIZE);
EXPECT_LE(s.rehashes, TABLE_SIZE_BITS);
//
// Check hash table distribution
//
auto ls = t->getLongStats(2);
EXPECT_EQ(ls.runs.size(), 2);
EXPECT_LT(ls.runs.rbegin()->first, 100); // Only look at second longest run
//
// Duplicate add of Handle
//
for (size_t i = 0; i < TABLE_SIZE; ++i) {
handles2.push_back(t->makeHandle(std::to_string(i)));
EXPECT_EQ(handles1[i], handles2[i]);
}
s = t->getStats();
EXPECT_EQ(s.size, TABLE_SIZE);
EXPECT_LE(s.rehashes, 0);
EXPECT_EQ(s.handles, 2*TABLE_SIZE);
//
// Now, delete each handle once. Basically nothing about the table should change
//
for (auto& h : handles1) { t->destroyHandle(h); }
s = t->getStats();
EXPECT_EQ(s.size, TABLE_SIZE);
EXPECT_EQ(s.handles, TABLE_SIZE);
EXPECT_EQ(s.maxSearch, 0);
EXPECT_EQ(s.rehashes, 0);
//
// Now empty the table
//
for (auto& h : handles2) { t->destroyHandle(h); }
s = t->getStats();
EXPECT_EQ(s.size, 0);
EXPECT_EQ(s.handles, 0);
EXPECT_GT(s.rehashes, TABLE_SIZE_BITS - 3); // Minimum table size
}

101
tst/unit/module_sim.cc Normal file
View File

@ -0,0 +1,101 @@
#undef NDEBUG
#include <assert.h>
#include <stdarg.h>
#include <signal.h>
#include <map>
#include <cstdlib>
#include <cstdio>
#include <ctime>
#include <cstdint>
#include <iostream>
#include <gtest/gtest.h>
#include <gmock/gmock.h>
#include "json/alloc.h"
#include "json/dom.h"
#include "json/stats.h"
#include "json/selector.h"
#include "module_sim.h"
//
// Simulate underlying zmalloc stuff, including malloc-size
//
static std::map<void *, size_t> malloc_sizes;
size_t malloced = 0;
std::string logtext;
static void *test_malloc(size_t s) {
void *ptr = malloc(s);
assert(malloc_sizes.find(ptr) == malloc_sizes.end());
malloc_sizes[ptr] = s;
malloced += s;
return ptr;
}
static size_t test_malloc_size(void *ptr) {
if (!ptr) return 0;
assert(malloc_sizes.find(ptr) != malloc_sizes.end());
return malloc_sizes[ptr];
}
static void test_free(void *ptr) {
if (!ptr) return;
assert(malloc_sizes.find(ptr) != malloc_sizes.end());
ASSERT_GE(malloced, malloc_sizes[ptr]);
malloced -= malloc_sizes[ptr];
malloc_sizes.erase(malloc_sizes.find(ptr));
free(ptr);
}
static void *test_realloc(void *old_ptr, size_t new_size) {
if (old_ptr == nullptr) return test_malloc(new_size);
assert(malloc_sizes.find(old_ptr) != malloc_sizes.end());
assert(malloced >= malloc_sizes[old_ptr]);
malloced -= malloc_sizes[old_ptr];
malloc_sizes.erase(malloc_sizes.find(old_ptr));
void *new_ptr = realloc(old_ptr, new_size);
assert(malloc_sizes.find(new_ptr) == malloc_sizes.end());
malloc_sizes[new_ptr] = new_size;
malloced += new_size;
return new_ptr;
}
std::string test_getLogText() {
std::string result = logtext;
logtext.resize(0);
return result;
}
static void test_log(ValkeyModuleCtx *ctx, const char *level, const char *fmt, ...) {
(void)ctx;
char buffer[256];
va_list arg;
va_start(arg, fmt);
int len = vsnprintf(buffer, sizeof(buffer), fmt, arg);
va_end(arg);
std::cerr << "Log(" << level << "): " << std::string(buffer, len) << "\n"; // make visible to ASSERT_EXIT
}
static void test__assert(const char *estr, const char *file, int line) {
ASSERT_TRUE(0) << "Assert(" << file << ":" << line << "): " << estr;
}
static long long test_Milliseconds() {
struct timespec t;
clock_gettime(CLOCK_REALTIME, &t);
return (t.tv_sec * 1000) + (t.tv_nsec / 1000000);
}
void setupValkeyModulePointers() {
ValkeyModule_Alloc = test_malloc;
ValkeyModule_Free = test_free;
ValkeyModule_Realloc = test_realloc;
ValkeyModule_MallocSize = test_malloc_size;
ValkeyModule_Log = test_log;
ValkeyModule__Assert = test__assert;
ValkeyModule_Strdup = strdup;
ValkeyModule_Milliseconds = test_Milliseconds;
memory_traps_control(true);
}

19
tst/unit/module_sim.h Normal file
View File

@ -0,0 +1,19 @@
//
// Simulate the Valkey Module Environment
//
#ifndef VALKEYJSONMODULE_TST_UNIT_MODULE_SIM_H_
#define VALKEYJSONMODULE_TST_UNIT_MODULE_SIM_H_
#include <cstddef>
#include <string>
extern size_t malloced; // Total currently allocated memory
void setupValkeyModulePointers();
std::string test_getLogText();
#endif // VALKEYJSONMODULE_TST_UNIT_MODULE_SIM_H_

1344
tst/unit/selector_test.cc Normal file
View File

File diff suppressed because it is too large Load Diff

32
tst/unit/stats_test.cc Normal file
View File

@ -0,0 +1,32 @@
#include <gtest/gtest.h>
#include "json/stats.h"
class StatsTest : public ::testing::Test {
};
TEST_F(StatsTest, testFindBucket) {
EXPECT_EQ(jsonstats_find_bucket(0), 0);
EXPECT_EQ(jsonstats_find_bucket(200), 0);
EXPECT_EQ(jsonstats_find_bucket(256), 1);
EXPECT_EQ(jsonstats_find_bucket(500), 1);
EXPECT_EQ(jsonstats_find_bucket(1024), 2);
EXPECT_EQ(jsonstats_find_bucket(2000), 2);
EXPECT_EQ(jsonstats_find_bucket(4*1024), 3);
EXPECT_EQ(jsonstats_find_bucket(5000), 3);
EXPECT_EQ(jsonstats_find_bucket(16*1024), 4);
EXPECT_EQ(jsonstats_find_bucket(50000), 4);
EXPECT_EQ(jsonstats_find_bucket(64*1024), 5);
EXPECT_EQ(jsonstats_find_bucket(100000), 5);
EXPECT_EQ(jsonstats_find_bucket(256*1024), 6);
EXPECT_EQ(jsonstats_find_bucket(1000000), 6);
EXPECT_EQ(jsonstats_find_bucket(1024*1024), 7);
EXPECT_EQ(jsonstats_find_bucket(4000000), 7);
EXPECT_EQ(jsonstats_find_bucket(4*1024*1024), 8);
EXPECT_EQ(jsonstats_find_bucket(5000000), 8);
EXPECT_EQ(jsonstats_find_bucket(16*1024*1024), 9);
EXPECT_EQ(jsonstats_find_bucket(20000000), 9);
EXPECT_EQ(jsonstats_find_bucket(60*1024*1024), 9);
EXPECT_EQ(jsonstats_find_bucket(64*1024*1024), 10);
EXPECT_EQ(jsonstats_find_bucket(90000000), 10);
EXPECT_EQ(jsonstats_find_bucket(1024*1024*1024), 10);
}

180
tst/unit/traps_test.cc Normal file
View File

@ -0,0 +1,180 @@
#undef NDEBUG
#include <assert.h>
#include <stdarg.h>
#include <signal.h>
#include <cstdlib>
#include <cstddef>
#include <cstring>
#include <cstdio>
#include <cstdint>
#include <cmath>
#include <memory>
#include <deque>
#include <vector>
#include <string>
#include <sstream>
#include <utility>
#include <iostream>
#include <unordered_map>
#include <map>
#include <set>
#include <gtest/gtest.h>
#include <gmock/gmock.h>
#include "json/alloc.h"
#include "json/dom.h"
#include "json/stats.h"
#include "json/selector.h"
#include "module_sim.h"
extern size_t hash_function(const char *, size_t);
/* Since unit tests run outside of Valkey server, we need to map Valkey'
* memory management functions to cstdlib functions. */
static void SetupAllocFuncs(size_t numShards) {
setupValkeyModulePointers();
//
// Now setup the KeyTable, the RapidJson library now depends on it
//
KeyTable::Config c;
c.malloc = memory_alloc;
c.free = memory_free;
c.hash = hash_function;
c.numShards = numShards;
keyTable = new KeyTable(c);
}
class TrapsTest : public ::testing::Test {
protected:
void SetUp() override {
JsonUtilCode rc = jsonstats_init();
ASSERT_EQ(rc, JSONUTIL_SUCCESS);
SetupAllocFuncs(16);
}
void TearDown() override {
delete keyTable;
keyTable = nullptr;
}
};
//
// See if we can startup and shutdown with no failures
//
TEST_F(TrapsTest, sanity) {
void *ptr = dom_alloc(15);
dom_free(ptr);
}
enum JTYPE {
JT_BOOLEAN,
JT_INTEGER,
JT_SHORT_STRING,
JT_LONG_STRING,
JT_SHORT_DOUBLE,
JT_LONG_DOUBLE,
JT_ARRAY,
JT_OBJECT,
JT_OBJECT_HT,
JT_NUM_TYPES
};
static void makeValue(JValue *v, JTYPE jt) {
std::string json;
switch (jt) {
case JT_BOOLEAN:
json = "true";
break;
case JT_INTEGER:
json = "1";
break;
case JT_SHORT_STRING:
json = "\"short\"";
break;
case JT_LONG_STRING:
json = "\"string of length large\"";
break;
case JT_SHORT_DOUBLE:
json = "1.2";
break;
case JT_LONG_DOUBLE:
json = "1.23456789101112";
break;
case JT_ARRAY:
json = "[1,2,3,4,5]";
break;
case JT_OBJECT:
json = "{\"a\":1}";
break;
case JT_OBJECT_HT:
json = "{";
for (auto s = 0; s < 1000; ++s) {
if (s != 0) json += ',';
json += '\"';
json += std::to_string(s);
json += "\":1";
}
json += '}';
break;
default:
ASSERT_TRUE(0);
}
JParser parser;
*v = parser.Parse(json.c_str(), json.length()).GetJValue();
}
//
// Test that keys properly honor corruption
//
TEST_F(TrapsTest, handle_corruption) {
for (auto corruption : {CORRUPT_PREFIX, CORRUPT_LENGTH, CORRUPT_SUFFIX}) {
for (auto jt : {JT_OBJECT, JT_OBJECT_HT}) {
JValue *v = new JValue;
makeValue(v, jt);
auto first = v->MemberBegin();
auto trap_pointer = &*(first->name);
memory_corrupt_memory(trap_pointer, corruption);
//
// Serialize this object
//
rapidjson::StringBuffer oss;
ASSERT_EXIT(dom_serialize_value(*v, nullptr, oss), testing::ExitedWithCode(1), "Validation Failure");
//
// Destruct it
//
ASSERT_EXIT(delete v, testing::ExitedWithCode(1), "Validation Failure");
//
// Cleanup
//
memory_uncorrupt_memory(trap_pointer, corruption);
delete v;
}
}
}
//
// Test out the JValue validate and dump functions
//
TEST_F(TrapsTest, jvalue_validation) {
std::string json =
"{ \"a\":1, \"b\":[1,2,\"this is a long string\",\"shortstr\",false,true,1.0,1.23456789012345,null]}";
JParser parser;
JValue *v = new JValue;
*v = parser.Parse(json.c_str(), json.length()).GetJValue();
std::ostringstream os;
DumpRedactedJValue(os, *v);
std::cerr << os.str() << "\n";
delete v;
}
//
// Test Log Stream
//
TEST_F(TrapsTest, test_log_stream) {
JValue v, v0;
v.SetArray();
v.PushBack(v0, allocator);
DumpRedactedJValue(v, nullptr, "level");
std::string log = test_getLogText();
std::cerr << log;
}

213
tst/unit/util_test.cc Normal file
View File

@ -0,0 +1,213 @@
#include <stdint.h>
#include <cstdlib>
#include <cstring>
#include <string>
#include <gtest/gtest.h>
#include "json/util.h"
#include "json/dom.h"
#include "json/alloc.h"
#include "json/stats.h"
#include "module_sim.h"
extern size_t dummy_malloc_size(void *);
class UtilTest : public ::testing::Test {
protected:
void SetUp() override {
JsonUtilCode rc = jsonstats_init();
ASSERT_EQ(rc, JSONUTIL_SUCCESS);
setupValkeyModulePointers();
}
};
TEST_F(UtilTest, testCodeToMessage) {
for (JsonUtilCode code=JSONUTIL_SUCCESS; code < JSONUTIL_LAST; code = JsonUtilCode(code + 1)) {
const char *msg = jsonutil_code_to_message(code);
EXPECT_TRUE(msg != nullptr);
if (code == JSONUTIL_SUCCESS || code == JSONUTIL_WRONG_NUM_ARGS ||
code == JSONUTIL_NX_XX_CONDITION_NOT_SATISFIED) {
EXPECT_STREQ(msg, "");
} else {
EXPECT_GT(strlen(msg), 0);
}
}
}
TEST_F(UtilTest, testDoubleToString) {
double v = 189.31;
char buf[BUF_SIZE_DOUBLE_JSON];
size_t len = jsonutil_double_to_string(v, buf, sizeof(buf));
EXPECT_STREQ(buf, "189.31");
EXPECT_EQ(len, strlen(buf));
}
TEST_F(UtilTest, testDoubleToStringRapidJson) {
double v = 189.31;
char buf[BUF_SIZE_DOUBLE_RAPID_JSON];
size_t len = jsonutil_double_to_string_rapidjson(v, buf, sizeof(buf));
EXPECT_STREQ(buf, "189.31");
EXPECT_EQ(len, strlen(buf));
}
TEST_F(UtilTest, testIsInt64) {
EXPECT_TRUE(jsonutil_is_int64(0));
EXPECT_TRUE(jsonutil_is_int64(1));
EXPECT_TRUE(jsonutil_is_int64(INT8_MAX));
EXPECT_TRUE(jsonutil_is_int64(INT8_MIN));
EXPECT_TRUE(jsonutil_is_int64(INT16_MAX));
EXPECT_TRUE(jsonutil_is_int64(INT16_MIN));
EXPECT_TRUE(jsonutil_is_int64(INT32_MAX));
EXPECT_TRUE(jsonutil_is_int64(INT32_MIN));
EXPECT_TRUE(jsonutil_is_int64(INT64_MAX >> 1));
EXPECT_TRUE(jsonutil_is_int64(8223372036854775807LL));
EXPECT_TRUE(jsonutil_is_int64(INT64_MIN));
EXPECT_FALSE(jsonutil_is_int64(1e28)); // out of range of int64
EXPECT_FALSE(jsonutil_is_int64(1.7e308)); // out of range of int64
EXPECT_FALSE(jsonutil_is_int64(-1e28)); // out of range of int64
EXPECT_FALSE(jsonutil_is_int64(-1.7e308)); // out of range of int64
EXPECT_TRUE(jsonutil_is_int64(108.0));
EXPECT_FALSE(jsonutil_is_int64(108.9));
EXPECT_FALSE(jsonutil_is_int64(108.0000001));
EXPECT_TRUE(jsonutil_is_int64(-108.0));
EXPECT_FALSE(jsonutil_is_int64(-108.9));
EXPECT_FALSE(jsonutil_is_int64(-108.0000001));
}
TEST_F(UtilTest, testMultiplyInt64_overflow) {
// should not overflow
int64_t res;
JsonUtilCode rc = jsonutil_multiply_int64(INT64_MAX, 1, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, INT64_MAX);
// should overflow
rc = jsonutil_multiply_int64(INT64_MAX, 2, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
rc = jsonutil_multiply_int64(INT64_MAX, INT64_MAX >> 1, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
rc = jsonutil_multiply_int64(INT64_MAX, INT64_MAX, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
}
TEST_F(UtilTest, testMultiplyInt64_overflow_negative) {
// should not overflow
int64_t res;
JsonUtilCode rc = jsonutil_multiply_int64(INT64_MIN, 1, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, INT64_MIN);
// should overflow
rc = jsonutil_multiply_int64(INT64_MIN, 2, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
rc = jsonutil_multiply_int64(INT64_MIN, INT64_MIN >> 1, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
rc = jsonutil_multiply_int64(INT64_MIN, INT64_MAX, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
rc = jsonutil_multiply_int64(INT64_MIN, INT64_MIN, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
}
TEST_F(UtilTest, testMultiplyDouble) {
double res;
JsonUtilCode rc = jsonutil_multiply_double(5e30, 2, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, 1e31);
rc = jsonutil_multiply_double(5.0e30, 2.0, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, 1.0e31);
}
TEST_F(UtilTest, testMultiplyDouble_overflow) {
// should not overflow
double res;
JsonUtilCode rc = jsonutil_multiply_double(1.7e308, 1.0, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, 1.7e308);
// should overflow
rc = jsonutil_multiply_double(1.7e308, 2.0, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
rc = jsonutil_multiply_double(1.7e308, 1.7e308, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
}
TEST_F(UtilTest, testMultiplyDouble_overflow_negative) {
// should not overflow
double res;
JsonUtilCode rc = jsonutil_multiply_double(-1.7e308, 1.0, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, -1.7e308);
// should overflow
rc = jsonutil_multiply_double(-1.7e308, 2.0, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
rc = jsonutil_multiply_double(-1.7e308, 1.7e308, &res);
EXPECT_EQ(rc, JSONUTIL_MULTIPLICATION_OVERFLOW);
}
TEST_F(UtilTest, testAddInt64_overflow) {
// should not overflow
int64_t res;
JsonUtilCode rc = jsonutil_add_int64(INT64_MAX, 0, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, INT64_MAX);
// should overflow
rc = jsonutil_add_int64(INT64_MAX, 1, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
rc = jsonutil_add_int64(INT64_MAX, INT64_MAX >> 1, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
rc = jsonutil_add_int64(INT64_MAX, INT64_MAX, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
}
TEST_F(UtilTest, testAddInt64_overflow_negative) {
// should not overflow
int64_t res;
JsonUtilCode rc = jsonutil_add_int64(INT64_MIN, 0, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, INT64_MIN);
// should overflow
rc = jsonutil_add_int64(INT64_MIN, -1, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
rc = jsonutil_add_int64(INT64_MIN, INT64_MIN >> 1, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
rc = jsonutil_add_int64(INT64_MIN, INT64_MIN, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
}
TEST_F(UtilTest, testAddDouble_overflow) {
// should not overflow
double res;
JsonUtilCode rc = jsonutil_add_double(1.7e308, 0.0, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, 1.7e308);
rc = jsonutil_add_double(1.7e308, 1.0, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, 1.7e308);
// should overflow
rc = jsonutil_add_double(1.7e308, 0.85e308, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
rc = jsonutil_add_double(1.7e308, 1.7e308, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
}
TEST_F(UtilTest, testAddDouble_overflow_negative) {
// should not overflow
double res;
JsonUtilCode rc = jsonutil_add_double(-1.7e308, 0.0, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, -1.7e308);
rc = jsonutil_add_double(-1.7e308, -1.0, &res);
EXPECT_EQ(rc, JSONUTIL_SUCCESS);
EXPECT_EQ(res, -1.7e308);
// should overflow
rc = jsonutil_add_double(-1.7e308, -0.85e308, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
rc = jsonutil_add_double(-1.7e308, -1.7e308, &res);
EXPECT_EQ(rc, JSONUTIL_ADDITION_OVERFLOW);
}