From e6407ed4057623f195be941b93e74cd70e6c9372 Mon Sep 17 00:00:00 2001
From: Artem Krylysov <artem.krylysov@datadoghq.com>
Date: Wed, 11 Aug 2021 08:43:44 -0400
Subject: [PATCH] Support building RocksDB with jemalloc (#542)

---
 .github/workflows/rust.yml | 6 ++++++
 Cargo.toml                 | 1 +
 librocksdb-sys/Cargo.toml  | 2 ++
 librocksdb-sys/build.rs    | 4 ++++
 4 files changed, 13 insertions(+)

diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 07311fb..20fe46d 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -99,3 +99,9 @@ jobs:
         with:
           command: test
           args: --features multi-threaded-cf
+      - name: Run rocksdb tests (jemalloc)
+        if: runner.os != 'Windows'
+        uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --features jemalloc
diff --git a/Cargo.toml b/Cargo.toml
index 7a5b880..9c8d487 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,6 +18,7 @@ exclude = [
 
 [features]
 default = ["snappy", "lz4", "zstd", "zlib", "bzip2"]
+jemalloc = ["librocksdb-sys/jemalloc"]
 valgrind = []
 snappy = ["librocksdb-sys/snappy"]
 lz4 = ["librocksdb-sys/lz4"]
diff --git a/librocksdb-sys/Cargo.toml b/librocksdb-sys/Cargo.toml
index 06f2012..e4d880e 100644
--- a/librocksdb-sys/Cargo.toml
+++ b/librocksdb-sys/Cargo.toml
@@ -13,6 +13,7 @@ links = "rocksdb"
 
 [features]
 default = [ "static" ]
+jemalloc = ["jemalloc-sys"]
 static = []
 snappy = []
 lz4 = []
@@ -22,6 +23,7 @@ bzip2 = []
 
 [dependencies]
 libc = "0.2"
+jemalloc-sys = { version = "0.3", features = ["unprefixed_malloc_on_supported_platforms"], optional = true }
 
 [dev-dependencies]
 const-cstr = "0.3"
diff --git a/librocksdb-sys/build.rs b/librocksdb-sys/build.rs
index e41065e..e018b4e 100644
--- a/librocksdb-sys/build.rs
+++ b/librocksdb-sys/build.rs
@@ -187,6 +187,10 @@ fn build_rocksdb() {
 
     config.define("ROCKSDB_SUPPORT_THREAD_LOCAL", None);
 
+    if cfg!(feature = "jemalloc") {
+        config.define("WITH_JEMALLOC", "ON");
+    }
+
     if target.contains("msvc") {
         config.flag("-EHsc");
     } else {