Move collection cache internals to PSRAM-aware allocators

src/esp_jsondb/collection/collection.cpp

@@ -3,6 +3,14 @@
 #include "../utils/fs_utils.h"
 #include "../utils/time_utils.h"
 
+namespace {
+std::shared_ptr<DocumentRecord> makeSharedDocumentRecord(bool usePSRAMBuffers) {
+	return std::allocate_shared<DocumentRecord>(
+	    JsonDbAllocator<DocumentRecord>(usePSRAMBuffers), usePSRAMBuffers
+	);
+}
+} // namespace
+
 Collection::Collection(
     ESPJsonDB &db,
     const std::string &name,
@@ -12,8 +20,10 @@ Collection::Collection(
     bool usePSRAMBuffers,
     fs::FS &fs
 )
-    : _db(&db), _name(name), _schema(schema), _baseDir(std::move(baseDir)), _cacheEnabled(true),
-      _usePSRAMBuffers(usePSRAMBuffers), _fs(&fs) {
+    : _db(&db), _name(name), _schema(schema),
+      _docs(std::less<std::string>{}, DocumentMapAllocator(usePSRAMBuffers)),
+      _baseDir(std::move(baseDir)), _cacheEnabled(true), _usePSRAMBuffers(usePSRAMBuffers),
+      _fs(&fs) {
 	(void)cacheEnabled;
 }
 
@@ -147,7 +157,7 @@ DbResult<std::string> Collection::create(JsonObjectConst data) {
 			recordStatus(res.status);
 			return res;
 		}
-		rec = std::make_shared<DocumentRecord>(_usePSRAMBuffers);
+		rec = makeSharedDocumentRecord(_usePSRAMBuffers);
 		rec->meta.createdAt = nowUtcMs();
 		rec->meta.updatedAt = rec->meta.createdAt;
 		rec->meta.id = ObjectId().toHex();
@@ -320,7 +330,7 @@ DbStatus Collection::updateOne(
 
 	// If not found and create requested, create a new record and apply mutator
 	if (!updated && create) {
-		auto rec = std::make_shared<DocumentRecord>(_usePSRAMBuffers);
+		auto rec = makeSharedDocumentRecord(_usePSRAMBuffers);
 		rec->meta.createdAt = nowUtcMs();
 		rec->meta.updatedAt = rec->meta.createdAt;
 		rec->meta.id = ObjectId().toHex();
@@ -409,7 +419,7 @@ DbStatus Collection::updateOne(const JsonDocument &filter, const JsonDocument &p
 
 	if (!updated && create) {
 		// Create a new document merging filter and patch
-		auto rec = std::make_shared<DocumentRecord>(_usePSRAMBuffers);
+		auto rec = makeSharedDocumentRecord(_usePSRAMBuffers);
 		rec->meta.createdAt = nowUtcMs();
 		rec->meta.updatedAt = rec->meta.createdAt;
 		rec->meta.id = ObjectId().toHex();
@@ -551,7 +561,7 @@ Collection::readDocFromFile(const std::string &baseDir, const std::string &id) {
 		recordStatus(res.status);
 		return res;
 	}
-	auto rec = std::make_shared<DocumentRecord>(_usePSRAMBuffers);
+	auto rec = makeSharedDocumentRecord(_usePSRAMBuffers);
 	rec->meta.id = id;
 	rec->meta.createdAt = nowUtcMs();
 	rec->meta.updatedAt = rec->meta.createdAt;
@@ -656,7 +666,7 @@ DbStatus Collection::updateOneNoCache(
 		}
 	}
 	if (create) {
-		auto rec = std::make_shared<DocumentRecord>(_usePSRAMBuffers);
+		auto rec = makeSharedDocumentRecord(_usePSRAMBuffers);
 		rec->meta.createdAt = nowUtcMs();
 		rec->meta.updatedAt = rec->meta.createdAt;
 		rec->meta.id = ObjectId().toHex();
@@ -730,7 +740,7 @@ DbStatus Collection::updateOneJsonNoCache(
 		return recordStatus(st);
 	}
 	if (create) {
-		auto rec = std::make_shared<DocumentRecord>(_usePSRAMBuffers);
+		auto rec = makeSharedDocumentRecord(_usePSRAMBuffers);
 		rec->meta.createdAt = nowUtcMs();
 		rec->meta.updatedAt = rec->meta.createdAt;
 		rec->meta.id = ObjectId().toHex();

src/esp_jsondb/collection/collection.h

@@ -132,11 +132,20 @@ class Collection {
 	}
 
   private:
+	using DocumentRecordPtr = std::shared_ptr<DocumentRecord>;
+	using DocumentMapValue = std::pair<const std::string, DocumentRecordPtr>;
+	using DocumentMapAllocator = JsonDbAllocator<DocumentMapValue>;
+	using DocumentMap = std::map<
+	    std::string,
+	    DocumentRecordPtr,
+	    std::less<std::string>,
+	    DocumentMapAllocator>;
+
 	ESPJsonDB *_db = nullptr;
 	std::string _name;
 	Schema _schema;
 	// Use shared_ptr to keep records alive while views exist
-	std::map<std::string, std::shared_ptr<DocumentRecord>> _docs;
+	DocumentMap _docs;
 	bool _dirty = false;
 	std::vector<std::string> _deletedIds; // files to remove on next flush
 	FrMutex _mu;                          // guards _docs, _deletedIds

src/esp_jsondb/utils/jsondb_allocator.h

@@ -10,6 +10,13 @@
 #define ESP_JSONDB_HAS_BUFFER_MANAGER 0
 #endif
 
+#if defined(ARDUINO_ARCH_ESP32) && __has_include(<esp_heap_caps.h>)
+#include <esp_heap_caps.h>
+#define ESP_JSONDB_HAS_ESP32_HEAP_CAPS 1
+#else
+#define ESP_JSONDB_HAS_ESP32_HEAP_CAPS 0
+#endif
+
 #include <cstddef>
 #include <cstdlib>
 #include <limits>
@@ -21,8 +28,16 @@ namespace jsondb_allocator_detail {
 inline void *allocate(std::size_t bytes, bool usePSRAMBuffers) noexcept {
 #if ESP_JSONDB_HAS_BUFFER_MANAGER
 	return ESPBufferManager::allocate(bytes, usePSRAMBuffers);
+#else
+#if ESP_JSONDB_HAS_ESP32_HEAP_CAPS
+	if (usePSRAMBuffers) {
+		if (void *psramMemory = heap_caps_malloc(bytes, MALLOC_CAP_SPIRAM)) {
+			return psramMemory;
+		}
+	}
 #else
 	(void)usePSRAMBuffers;
+#endif
 	return std::malloc(bytes);
 #endif
 }
@@ -38,8 +53,16 @@ inline void deallocate(void *ptr) noexcept {
 inline void *reallocate(void *ptr, std::size_t bytes, bool usePSRAMBuffers) noexcept {
 #if ESP_JSONDB_HAS_BUFFER_MANAGER
 	return ESPBufferManager::reallocate(ptr, bytes, usePSRAMBuffers);
+#else
+#if ESP_JSONDB_HAS_ESP32_HEAP_CAPS
+	if (usePSRAMBuffers) {
+		if (void *psramMemory = heap_caps_realloc(ptr, bytes, MALLOC_CAP_SPIRAM)) {
+			return psramMemory;
+		}
+	}
 #else
 	(void)usePSRAMBuffers;
+#endif
 	return std::realloc(ptr, bytes);
 #endif
 }

test/dbTest.cpp

@@ -59,6 +59,7 @@ void DbTester::run() {
 	delayedCollectionSyncNowFallbackTest();
 	delayedCollectionDropBeforeLoadTest();
 	delayedCollectionConfigNormalizationTest();
+	psramBufferWiringTest();
 	printDBDiag();
 	teardownLifecycle();
 }

test/dbTest.h

@@ -47,6 +47,7 @@ class DbTester {
 	void delayedCollectionSyncNowFallbackTest();
 	void delayedCollectionDropBeforeLoadTest();
 	void delayedCollectionConfigNormalizationTest();
+	void psramBufferWiringTest();
 	// Utils
 	void printDBDiag();
 	void teardownLifecycle();

test/psramTests.cpp

@@ -0,0 +1,115 @@
+#include "dbTest.h"
+
+void DbTester::psramBufferWiringTest() {
+	auto clearStatus = db.dropAll();
+	if (!clearStatus.ok()) {
+		ESP_LOGE(DB_TESTER_TAG, "dropAll() failed before PSRAM wiring test: %s", clearStatus.message);
+		return;
+	}
+
+	db.deinit();
+
+	ESPJsonDBConfig cfg;
+	cfg.autosync = false;
+	cfg.usePSRAMBuffers = true;
+
+	auto initStatus = db.init("/test_db", cfg);
+	if (!initStatus.ok()) {
+		ESP_LOGE(DB_TESTER_TAG, "init(usePSRAMBuffers=true) failed: %s", initStatus.message);
+		return;
+	}
+
+	JsonDocument seedDoc;
+	seedDoc["kind"] = "seed";
+	auto createRes = db.create("psram_wiring", seedDoc.as<JsonObjectConst>());
+	if (!createRes.status.ok()) {
+		ESP_LOGE(DB_TESTER_TAG, "create(seed) failed in PSRAM wiring test: %s", createRes.status.message);
+		return;
+	}
+
+	auto lambdaUpsert = db.updateOne(
+	    "psram_wiring",
+	    [](const DocView &v) {
+		    return v["kind"] == "lambda_match";
+	    },
+	    [](DocView &v) {
+		    v["kind"].set("lambda_created");
+	    },
+	    true
+	);
+	if (!lambdaUpsert.ok()) {
+		ESP_LOGE(
+		    DB_TESTER_TAG,
+		    "updateOne(lambda, create=true) failed in PSRAM wiring test: %s",
+		    lambdaUpsert.message
+		);
+		return;
+	}
+
+	JsonDocument filterDoc;
+	filterDoc["kind"] = "json_match";
+	JsonDocument patchDoc;
+	patchDoc["kind"] = "json_created";
+	patchDoc["marker"] = true;
+	auto jsonUpsert = db.updateOne("psram_wiring", filterDoc, patchDoc, true);
+	if (!jsonUpsert.ok()) {
+		ESP_LOGE(
+		    DB_TESTER_TAG,
+		    "updateOne(json, create=true) failed in PSRAM wiring test: %s",
+		    jsonUpsert.message
+		);
+		return;
+	}
+
+	auto syncStatus = db.syncNow();
+	if (!syncStatus.ok()) {
+		ESP_LOGE(DB_TESTER_TAG, "syncNow() failed in PSRAM wiring test: %s", syncStatus.message);
+		return;
+	}
+
+	db.deinit();
+	initStatus = db.init("/test_db", cfg);
+	if (!initStatus.ok()) {
+		ESP_LOGE(DB_TESTER_TAG, "re-init(usePSRAMBuffers=true) failed: %s", initStatus.message);
+		return;
+	}
+
+	auto findRes = db.findMany("psram_wiring", [](const DocView &) {
+		return true;
+	});
+	if (!findRes.status.ok()) {
+		ESP_LOGE(DB_TESTER_TAG, "findMany(psram_wiring) failed: %s", findRes.status.message);
+		return;
+	}
+	if (findRes.value.size() != 3) {
+		ESP_LOGE(
+		    DB_TESTER_TAG,
+		    "PSRAM wiring preload expected 3 docs, got %d",
+		    static_cast<int>(findRes.value.size())
+		);
+		return;
+	}
+
+	bool hasSeed = false;
+	bool hasLambdaCreated = false;
+	bool hasJsonCreated = false;
+	for (const auto &doc : findRes.value) {
+		const char *kind = doc["kind"].as<const char *>();
+		if (!kind)
+			continue;
+		if (strcmp(kind, "seed") == 0) {
+			hasSeed = true;
+		} else if (strcmp(kind, "lambda_created") == 0) {
+			hasLambdaCreated = true;
+		} else if (strcmp(kind, "json_created") == 0) {
+			hasJsonCreated = true;
+		}
+	}
+
+	if (!hasSeed || !hasLambdaCreated || !hasJsonCreated) {
+		ESP_LOGE(DB_TESTER_TAG, "PSRAM wiring test documents mismatch after preload");
+		return;
+	}
+
+	ESP_LOGI(DB_TESTER_TAG, "PSRAM buffer wiring test passed");
+}