chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Provide explicit argument for asset cache ID.

Browse Source 
This will allow assets to cache compiled results with different parameters. For example, we could cache both an RGB version of a texture and a Mono version.

PiperOrigin-RevId: 811341377
Change-Id: Ifd6a4ff88661a2b38be47c02d6a666de2c202885
This commit is contained in:
Sam Haves
2025-09-25 07:23:11 -07:00
committed by Copybara-Service
parent 0baac58999
commit 922aebbeed
5 changed files with 79 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/user/user_cache.cc
+6 -6
View File
@@ -64,17 +64,17 @@ const std::string* mjCCache::HasAsset(const std::string& id) {
// inserts an asset into the cache, if asset is already in the cache, its data // inserts an asset into the cache, if asset is already in the cache, its data
// is updated only if the timestamps disagree // is updated only if the timestamps disagree
bool mjCCache::Insert(const std::string& modelname, const mjResource *resource, bool mjCCache::Insert(const std::string& modelname, const std::string& id, const mjResource *resource,
std::shared_ptr<const void> data, std::size_t size) { std::shared_ptr<const void> data, std::size_t size) {
std::lock_guard<std::mutex> lock(mutex_); std::lock_guard<std::mutex> lock(mutex_);
// check if asset is too large to fit in the cache // check if asset is too large to fit in the cache
if ((size_ + size > capacity_) && if ((size_ + size > capacity_) &&
lookup_.find(resource->name) == lookup_.end()) { lookup_.find(id) == lookup_.end()) {
return false; return false;
} }
mjCAsset asset(modelname, resource, data, size); mjCAsset asset(modelname, id, resource, data, size);
auto [it, inserted] = lookup_.insert({resource->name, asset}); auto [it, inserted] = lookup_.insert({id, asset});
mjCAsset* asset_ptr = &(it->second); mjCAsset* asset_ptr = &(it->second);
if (!inserted) { if (!inserted) {
@@ -104,9 +104,9 @@ bool mjCCache::Insert(const std::string& modelname, const mjResource *resource,
// populate data from the cache into the given function, return true if data was // populate data from the cache into the given function, return true if data was
// copied // copied
bool mjCCache::PopulateData(const mjResource* resource, mjCDataFunc fn) { bool mjCCache::PopulateData(const std::string& id, const mjResource* resource, mjCDataFunc fn) {
std::lock_guard<std::mutex> lock(mutex_); std::lock_guard<std::mutex> lock(mutex_);
auto it = lookup_.find(resource->name); auto it = lookup_.find(id);
if (it == lookup_.end()) { if (it == lookup_.end()) {
return false; return false;
} }
src/user/user_cache.h
+4 -4
View File
@@ -39,9 +39,9 @@ typedef void (*mjCDeallocFunc)(const void*);
class mjCAsset { class mjCAsset {
friend class mjCCache; friend class mjCCache;
public: public:
mjCAsset(std::string modelname, const mjResource* resource, mjCAsset(std::string modelname, std::string id, const mjResource* resource,
std::shared_ptr<const void> data, std::size_t size) : std::shared_ptr<const void> data, std::size_t size) :
id_(resource->name), timestamp_(resource->timestamp), id_(id), timestamp_(resource->timestamp),
size_(size), data_(std::move(data)) { size_(size), data_(std::move(data)) {
AddReference(modelname); AddReference(modelname);
} }
@@ -136,11 +136,11 @@ class mjCCache {
// inserts an asset into the cache, if asset is already in the cache, its data // inserts an asset into the cache, if asset is already in the cache, its data
// is updated only if the timestamps disagree // is updated only if the timestamps disagree
bool Insert(const std::string& modelname, const mjResource *resource, bool Insert(const std::string& modelname, const std::string& id, const mjResource *resource,
std::shared_ptr<const void> data, std::size_t size); std::shared_ptr<const void> data, std::size_t size);
// populate data from the cache into the given function // populate data from the cache into the given function
bool PopulateData(const mjResource* resource, mjCDataFunc fn); bool PopulateData(const std::string& id, const mjResource* resource, mjCDataFunc fn);
// deletes the asset from the cache with the given id // deletes the asset from the cache with the given id
void DeleteAsset(const std::string& id); void DeleteAsset(const std::string& id);
src/user/user_mesh.cc
+2 -2
View File
@@ -513,7 +513,7 @@ void mjCMesh::CacheMesh(mjCCache* cache, const mjResource* resource) {
const mjCMesh* mesh = static_cast<const mjCMesh*>(data); const mjCMesh* mesh = static_cast<const mjCMesh*>(data);
delete mesh; delete mesh;
}); });
cache->Insert("", resource, cached_data, size); cache->Insert("", resource->name, resource, cached_data, size);
} }
namespace { namespace {
@@ -1231,7 +1231,7 @@ bool mjCMesh::LoadCachedMesh(mjCCache *cache, const mjResource* resource) {
}; };
// check that cached asset has all data // check that cached asset has all data
return cache->PopulateData(resource, process_mesh); return cache->PopulateData(resource->name, resource, process_mesh);
} }
src/user/user_objects.cc
+2 -2
View File
@@ -91,7 +91,7 @@ PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource,
}; };
// try loading from cache // try loading from cache
if (cache && cache->PopulateData(resource, callback)) { if (cache && cache->PopulateData(resource->name, resource, callback)) {
return image; return image;
} }
@@ -139,7 +139,7 @@ PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource,
std::shared_ptr<const void> cached_data(cached_image, +[] (const void* data) { std::shared_ptr<const void> cached_data(cached_image, +[] (const void* data) {
delete static_cast<const PNGImage*>(data); delete static_cast<const PNGImage*>(data);
}); });
cache->Insert("", resource, cached_data, size); cache->Insert("", resource->name, resource, cached_data, size);
} }
return image; return image;
test/user/user_cache_test.cc
+65 -63
View File
@@ -39,30 +39,31 @@ namespace {
constexpr int kMaxSize = 100; // in bytes constexpr int kMaxSize = 100; // in bytes
constexpr std::string kText = "Hello World"; constexpr std::string kText = "Hello World";
constexpr std::string kTestId = "Test";
constexpr std::string kModel = "myModel"; constexpr std::string kModel = "myModel";
constexpr std::string kFile = "hello.txt"; constexpr std::string kFile = "hello.txt";
void CacheText(mjCCache& cache, const std::string& model, void CacheText(mjCCache& cache, const std::string& id, const std::string& model,
const std::string& name, const std::string& text) { const std::string& name, const std::string& text) {
mjVFS vfs; mjVFS vfs;
mj_defaultVFS(&vfs); mj_defaultVFS(&vfs);
mj_addBufferVFS(&vfs, name.c_str(), text.data(), text.size()); mj_addBufferVFS(&vfs, name.c_str(), text.data(), text.size());
mjResource* resource = mju_openResource("", name.c_str(), &vfs, nullptr, 0); mjResource* resource = mju_openResource("", name.c_str(), &vfs, nullptr, 0);
std::shared_ptr<const void> data(&text, +[](const void* data) {}); std::shared_ptr<const void> data(&text, +[](const void* data) {});
cache.Insert(model, resource, data, text.size()); cache.Insert(model, id, resource, data, text.size());
mju_closeResource(resource); mju_closeResource(resource);
mj_deleteVFS(&vfs); mj_deleteVFS(&vfs);
} }
std::optional<std::string> std::optional<std::string>
GetCachedText(mjCCache& cache, const std::string& model, GetCachedText(mjCCache& cache, const std::string& id, const std::string& model,
const std::string& name, const std::string& text) { const std::string& name, const std::string& text) {
std::string cached_text; std::string cached_text;
mjVFS vfs; mjVFS vfs;
mj_defaultVFS(&vfs); mj_defaultVFS(&vfs);
mj_addBufferVFS(&vfs, name.c_str(), text.data(), std::strlen(text.c_str())); mj_addBufferVFS(&vfs, name.c_str(), text.data(), std::strlen(text.c_str()));
mjResource* resource = mju_openResource("", name.c_str(), &vfs, nullptr, 0); mjResource* resource = mju_openResource("", name.c_str(), &vfs, nullptr, 0);
bool inserted = cache.PopulateData(resource, bool inserted = cache.PopulateData(id, resource,
[&cached_text](const void* data) { [&cached_text](const void* data) {
cached_text = *(static_cast<const std::string*>(data)); cached_text = *(static_cast<const std::string*>(data));
return true; return true;
@@ -80,25 +81,26 @@ TEST(CacheTest, SizeTest) {
TEST(CacheTest, InsertSuccess) { TEST(CacheTest, InsertSuccess) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, kModel, kFile, kText); CacheText(cache, kTestId, kModel, kFile, kText);
auto cached_text = GetCachedText(cache, kModel, kFile, kText); auto cached_text = GetCachedText(cache, kTestId, kModel, kFile, kText);
EXPECT_THAT(cached_text.value(), StrEq(kText)); EXPECT_THAT(cached_text.value(), StrEq(kText));
} }
TEST(CacheTest, InsertFailure) { TEST(CacheTest, InsertFailure) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, kModel, kFile, kText); CacheText(cache, kTestId, kModel, kFile, kText);
auto cached_text = GetCachedText(cache, kModel, "hello2.txt", kText); auto cached_text = GetCachedText(
cache, "WrongId", kModel, "hello2.txt", kText);
EXPECT_EQ(cached_text, std::nullopt); EXPECT_EQ(cached_text, std::nullopt);
} }
TEST(CacheTest, InsertReplace) { TEST(CacheTest, InsertReplace) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
const std::string kUpdatedText = "Goodbye World"; const std::string kUpdatedText = "Goodbye World";
CacheText(cache, kModel, kFile, kText); CacheText(cache, kTestId, kModel, kFile, kText);
CacheText(cache, kModel, kFile, kUpdatedText); CacheText(cache, kTestId, kModel, kFile, kUpdatedText);
auto cached_text = GetCachedText(cache, kModel, kFile, kUpdatedText); auto cached_text = GetCachedText(cache, kTestId, kModel, kFile, kUpdatedText);
EXPECT_THAT(cached_text.value(), StrEq(kUpdatedText)); EXPECT_THAT(cached_text.value(), StrEq(kUpdatedText));
} }
@@ -108,146 +110,146 @@ TEST(CacheTest, Limit1) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
EXPECT_THAT(cache.Capacity(), kMaxSize); EXPECT_THAT(cache.Capacity(), kMaxSize);
CacheText(cache, "fil.xml", "foo.obj", kText); CacheText(cache, "foo", "fil.xml", "foo.obj", kText);
CacheText(cache, "fil.xml", "bar.obj", kText); CacheText(cache, "bar", "fil.xml", "bar.obj", kText);
// access asset foo twice, bar one // access asset foo twice, bar one
GetCachedText(cache, "file.xml", "foo.obj", kText); GetCachedText(cache, "foo", "file.xml", "foo.obj", kText);
GetCachedText(cache, "file.xml", "foo.obj", kText); GetCachedText(cache, "foo", "file.xml", "foo.obj", kText);
GetCachedText(cache, "file.xml", "bar.obj", kText); GetCachedText(cache, "bar", "file.xml", "bar.obj", kText);
// make max size so cache can hold only one asset // make max size so cache can hold only one asset
cache.SetCapacity(12); cache.SetCapacity(12);
// foo should still be in cache // foo should still be in cache
EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("foo"), NotNull());
// bar was accessed less, so is removed // bar was accessed less, so is removed
EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("bar"), IsNull());
} }
// Trim cache based off of insert order // Trim cache based off of insert order
TEST(CacheTest, Limit2) { TEST(CacheTest, Limit2) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, "file.xml", "foo.obj", kText); CacheText(cache, "foo", "file.xml", "foo.obj", kText);
CacheText(cache, "file.xml", "bar.obj", kText); CacheText(cache, "bar", "file.xml", "bar.obj", kText);
// get each asset once // get each asset once
GetCachedText(cache, "file.xml", "foo.obj", kText); GetCachedText(cache, "foo", "file.xml", "foo.obj", kText);
GetCachedText(cache, "file.xml", "bar.obj", kText); GetCachedText(cache, "bar", "file.xml", "bar.obj", kText);
// make max size so cache can hold only one asset // make max size so cache can hold only one asset
cache.SetCapacity(12); cache.SetCapacity(12);
// foo should be gone because it's older // foo should be gone because it's older
EXPECT_THAT(cache.HasAsset("foo.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("foo"), IsNull());
// bar should still be in cache // bar should still be in cache
EXPECT_THAT(cache.HasAsset("bar.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("bar"), NotNull());
} }
// stress test with large asset // stress test with large asset
TEST(CacheTest, Limit3) { TEST(CacheTest, Limit3) {
mjCCache cache(12); mjCCache cache(12);
CacheText(cache, "file.xml", "foo.obj", kText); CacheText(cache, "foo", "file.xml", "foo.obj", kText);
CacheText(cache, "file.xml", "bar.obj", kText); CacheText(cache, "bar", "file.xml", "bar.obj", kText);
// foo should still be in cache // foo should still be in cache
EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("foo"), NotNull());
// bar could not be inserted because it's too large // bar could not be inserted because it's too large
EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("bar"), IsNull());
} }
TEST(CacheTest, Limit4) { TEST(CacheTest, Limit4) {
mjCCache cache(12); mjCCache cache(12);
CacheText(cache, "file.xml", "foo.obj", kText); CacheText(cache, "foo", "file.xml", "foo.obj", kText);
CacheText(cache, "file.xml", "bar.obj", kText); CacheText(cache, "bar", "file.xml", "bar.obj", kText);
EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("foo"), NotNull());
// cache is full, so bar can't be inserted // cache is full, so bar can't be inserted
EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("bar"), IsNull());
} }
TEST(CacheTest, ResetAll) { TEST(CacheTest, ResetAll) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, "file1.xml", "foo.obj", kText); CacheText(cache, "foo", "file1.xml", "foo.obj", kText);
CacheText(cache, "file2.xml", "bar.obj", kText); CacheText(cache, "bar", "file2.xml", "bar.obj", kText);
EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("foo"), NotNull());
EXPECT_THAT(cache.HasAsset("bar.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("bar"), NotNull());
cache.Reset(); cache.Reset();
EXPECT_THAT(cache.HasAsset("foo.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("foo"), IsNull());
EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("bar"), IsNull());
} }
TEST(CacheTest, ResetModel1) { TEST(CacheTest, ResetModel1) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, "file1.xml", "foo.obj", kText); CacheText(cache, "foo", "file1.xml", "foo.obj", kText);
CacheText(cache, "file1.xml", "bar.obj", kText); CacheText(cache, "bar", "file1.xml", "bar.obj", kText);
CacheText(cache, "file2.xml", "bar.obj", kText); CacheText(cache, "bar", "file2.xml", "bar.obj", kText);
cache.Reset("file2.xml"); cache.Reset("file2.xml");
EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("foo"), NotNull());
EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("bar"), IsNull());
} }
TEST(CacheTest, ResetModel2) { TEST(CacheTest, ResetModel2) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, "file1.xml", "foo.obj", kText); CacheText(cache, "foo", "file1.xml", "foo.obj", kText);
CacheText(cache, "file2.xml", "foo.obj", kText); CacheText(cache, "foo", "file2.xml", "foo.obj", kText);
CacheText(cache, "file2.xml", "bar.obj", kText); CacheText(cache, "bar", "file2.xml", "bar.obj", kText);
cache.Reset("file2.xml"); cache.Reset("file2.xml");
EXPECT_THAT(cache.HasAsset("foo.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("foo"), IsNull());
EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("bar"), IsNull());
} }
TEST(CacheTest, RemoveModel1) { TEST(CacheTest, RemoveModel1) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, "file1.xml", "foo.obj", kText); CacheText(cache, "foo", "file1.xml", "foo.obj", kText);
CacheText(cache, "file1.xml", "bar.obj", kText); CacheText(cache, "bar", "file1.xml", "bar.obj", kText);
CacheText(cache, "file2.xml", "bar.obj", kText); CacheText(cache, "bar", "file2.xml", "bar.obj", kText);
cache.RemoveModel("file2.xml"); cache.RemoveModel("file2.xml");
EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("foo"), NotNull());
EXPECT_THAT(cache.HasAsset("bar.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("bar"), NotNull());
} }
TEST(CacheTest, RemoveModel2) { TEST(CacheTest, RemoveModel2) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, "file1.xml", "foo.obj", kText); CacheText(cache, "foo", "file1.xml", "foo.obj", kText);
CacheText(cache, "file2.xml", "bar.obj", kText); CacheText(cache, "bar", "file2.xml", "bar.obj", kText);
cache.Reset("file2.xml"); cache.Reset("file2.xml");
EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("foo"), NotNull());
EXPECT_THAT(cache.HasAsset("bar.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("bar"), IsNull());
} }
TEST(CacheTest, DeleteAssetSuccess) { TEST(CacheTest, DeleteAssetSuccess) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, "file.xml", "foo.obj", kText); CacheText(cache, "foo", "file.xml", "foo.obj", kText);
cache.DeleteAsset("foo.obj"); cache.DeleteAsset("foo");
EXPECT_THAT(cache.HasAsset("foo.obj"), IsNull()); EXPECT_THAT(cache.HasAsset("foo"), IsNull());
} }
TEST(CacheTest, DeleteAssetFailure) { TEST(CacheTest, DeleteAssetFailure) {
mjCCache cache(kMaxSize); mjCCache cache(kMaxSize);
CacheText(cache, "file.xml", "foo.obj", kText); CacheText(cache, "foo", "file.xml", "foo.obj", kText);
cache.DeleteAsset("bar.obj"); cache.DeleteAsset("bar");
EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull()); EXPECT_THAT(cache.HasAsset("foo"), NotNull());
} }
} // namespace } // namespace