SharedAssetDepot.h

#pragma once
 
#include <CesiumAsync/AsyncSystem.h>
#include <CesiumAsync/Future.h>
#include <CesiumAsync/IAssetAccessor.h>
#include <CesiumUtility/DoublyLinkedList.h>
#include <CesiumUtility/IDepotOwningAsset.h>
#include <CesiumUtility/IntrusivePointer.h>
#include <CesiumUtility/ReferenceCounted.h>
#include <CesiumUtility/Result.h>
 
#include <cstddef>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <string>
#include <unordered_map>
 
namespace CesiumUtility {
template <typename T> class SharedAsset;
}
 
namespace CesiumAsync {
 
template <typename TAssetType, typename TAssetKey>
class CESIUMASYNC_API SharedAssetDepot
    : public CesiumUtility::ReferenceCountedThreadSafe<
          SharedAssetDepot<TAssetType, TAssetKey>>,
      public CesiumUtility::IDepotOwningAsset<TAssetType> {
public:
  int64_t inactiveAssetSizeLimitBytes = 16 * 1024 * 1024;
 
  using FactorySignature =
      CesiumAsync::Future<CesiumUtility::ResultPointer<TAssetType>>(
          const AsyncSystem& asyncSystem,
          const std::shared_ptr<IAssetAccessor>& pAssetAccessor,
          const TAssetKey& key);
 
  SharedAssetDepot(std::function<FactorySignature> factory)
      : _assets(),
        _assetsByPointer(),
        _deletionCandidates(),
        _totalDeletionCandidateMemoryUsage(0),
        _mutex(),
        _factory(std::move(factory)),
        _pKeepAlive(nullptr) {}
 
  virtual ~SharedAssetDepot() {
    // Ideally, when the depot is destroyed, all the assets it owns would become
    // independent assets. But this is extremely difficult to manage in a
    // thread-safe manner.
 
    // Since we're in the destructor, we can be sure no one has a reference to
    // this instance anymore. That means that no other thread can be executing
    // `getOrCreate`, and no async asset creations are in progress.
 
    // However, if assets owned by this depot are still alive, then other
    // threads can still be calling addReference / releaseReference on some of
    // our assets even while we're running the depot's destructor. Which means
    // that we can end up in `markDeletionCandidate` at the same time the
    // destructor is running. And in fact it's possible for a `SharedAsset` with
    // especially poor timing to call into a `SharedAssetDepot` just after it is
    // destroyed.
 
    // To avoid this, we use the _pKeepAlive field to maintain an artificial
    // reference to this depot whenever it owns live assets. This should keep
    // this destructor from being called except when all of its assets are also
    // in the _deletionCandidates list.
 
    CESIUM_ASSERT(this->_assets.size() == this->_deletionCandidates.size());
  }
 
  SharedFuture<CesiumUtility::ResultPointer<TAssetType>> getOrCreate(
      const AsyncSystem& asyncSystem,
      const std::shared_ptr<IAssetAccessor>& pAssetAccessor,
      const TAssetKey& assetKey) {
    // We need to take care here to avoid two assets starting to load before the
    // first asset has added an entry and set its maybePendingAsset field.
    std::unique_lock lock(this->_mutex);
 
    auto existingIt = this->_assets.find(assetKey);
    if (existingIt != this->_assets.end()) {
      // We've already loaded (or are loading) an asset with this ID - we can
      // just use that.
      const AssetEntry& entry = *existingIt->second;
      if (entry.maybePendingAsset) {
        // Asset is currently loading.
        return *entry.maybePendingAsset;
      } else {
        return asyncSystem.createResolvedFuture(entry.toResultUnderLock())
            .share();
      }
    }
 
    // Calling the factory function while holding the mutex unnecessarily
    // limits parallelism. It can even lead to a bug in the scenario where the
    // `thenInWorkerThread` continuation is invoked immediately in the current
    // thread, before `thenInWorkerThread` itself returns. That would result
    // in an attempt to lock the mutex recursively, which is not allowed.
 
    // So we jump through some hoops here to publish "this thread is working
    // on it", then unlock the mutex, and _then_ actually call the factory
    // function.
    Promise<void> promise = asyncSystem.createPromise<void>();
 
    // We haven't loaded or started to load this asset yet.
    // Let's do that now.
    CesiumUtility::IntrusivePointer<SharedAssetDepot<TAssetType, TAssetKey>>
        pDepot = this;
    CesiumUtility::IntrusivePointer<AssetEntry> pEntry =
        new AssetEntry(assetKey);
 
    auto future =
        promise.getFuture()
            .thenImmediately([pDepot, pEntry, asyncSystem, pAssetAccessor]() {
              return pDepot->_factory(asyncSystem, pAssetAccessor, pEntry->key);
            })
            .catchImmediately([](std::exception&& e) {
              return CesiumUtility::Result<
                  CesiumUtility::IntrusivePointer<TAssetType>>(
                  CesiumUtility::ErrorList::error(
                      std::string("Error creating asset: ") + e.what()));
            })
            .thenInWorkerThread(
                [pDepot, pEntry](
                    CesiumUtility::Result<
                        CesiumUtility::IntrusivePointer<TAssetType>>&& result) {
                  std::lock_guard lock(pDepot->_mutex);
 
                  if (result.pValue) {
                    result.pValue->_pDepot = pDepot.get();
                    pDepot->_assetsByPointer[result.pValue.get()] =
                        pEntry.get();
                  }
 
                  // Now that this asset is owned by the depot, we exclusively
                  // control its lifetime with a std::unique_ptr.
                  pEntry->pAsset =
                      std::unique_ptr<TAssetType>(result.pValue.get());
                  pEntry->errorsAndWarnings = std::move(result.errors);
                  pEntry->maybePendingAsset.reset();
 
                  // The asset is initially live because we have an
                  // IntrusivePointer to it right here. So make sure the depot
                  // stays alive, too.
                  pDepot->_pKeepAlive = pDepot;
 
                  return pEntry->toResultUnderLock();
                });
 
    SharedFuture<CesiumUtility::ResultPointer<TAssetType>> sharedFuture =
        std::move(future).share();
 
    pEntry->maybePendingAsset = sharedFuture;
 
    auto [it, added] = this->_assets.emplace(assetKey, pEntry);
 
    // Should always be added successfully, because we checked above that the
    // asset key doesn't exist in the map yet.
    CESIUM_ASSERT(added);
 
    // Unlock the mutex and then call the factory function.
    lock.unlock();
    promise.resolve();
 
    return sharedFuture;
  }
 
  size_t getAssetCount() const {
    std::lock_guard lock(this->_mutex);
    return this->_assets.size();
  }
 
  size_t getActiveAssetCount() const {
    std::lock_guard lock(this->_mutex);
    return this->_assets.size() - this->_deletionCandidates.size();
  }
 
  size_t getInactiveAssetCount() const {
    std::lock_guard lock(this->_mutex);
    return this->_deletionCandidates.size();
  }
 
  int64_t getInactiveAssetTotalSizeBytes() const {
    std::lock_guard lock(this->_mutex);
    return this->_totalDeletionCandidateMemoryUsage;
  }
 
private:
  // Disable copy
  void operator=(const SharedAssetDepot<TAssetType, TAssetKey>& other) = delete;
 
  void markDeletionCandidate(const TAssetType& asset, bool threadOwnsDepotLock)
      override {
    if (threadOwnsDepotLock) {
      this->markDeletionCandidateUnderLock(asset);
    } else {
      std::lock_guard lock(this->_mutex);
      this->markDeletionCandidateUnderLock(asset);
    }
  }
 
  void markDeletionCandidateUnderLock(const TAssetType& asset) {
    auto it = this->_assetsByPointer.find(const_cast<TAssetType*>(&asset));
    CESIUM_ASSERT(it != this->_assetsByPointer.end());
    if (it == this->_assetsByPointer.end()) {
      return;
    }
 
    CESIUM_ASSERT(it->second != nullptr);
 
    AssetEntry& entry = *it->second;
    entry.sizeInDeletionList = asset.getSizeBytes();
    this->_totalDeletionCandidateMemoryUsage += entry.sizeInDeletionList;
 
    this->_deletionCandidates.insertAtTail(entry);
 
    if (this->_totalDeletionCandidateMemoryUsage >
        this->inactiveAssetSizeLimitBytes) {
      // Delete the deletion candidates until we're below the limit.
      while (this->_deletionCandidates.size() > 0 &&
             this->_totalDeletionCandidateMemoryUsage >
                 this->inactiveAssetSizeLimitBytes) {
        AssetEntry* pOldEntry = this->_deletionCandidates.head();
        this->_deletionCandidates.remove(*pOldEntry);
 
        this->_totalDeletionCandidateMemoryUsage -=
            pOldEntry->sizeInDeletionList;
 
        CESIUM_ASSERT(
            pOldEntry->pAsset == nullptr ||
            pOldEntry->pAsset->_referenceCount == 0);
 
        if (pOldEntry->pAsset) {
          this->_assetsByPointer.erase(pOldEntry->pAsset.get());
        }
 
        // This will actually delete the asset.
        this->_assets.erase(pOldEntry->key);
      }
    }
 
    // If this depot is not managing any live assets, then we no longer need to
    // keep it alive.
    if (this->_assets.size() == this->_deletionCandidates.size()) {
      this->_pKeepAlive.reset();
    }
  }
 
  void unmarkDeletionCandidate(
      const TAssetType& asset,
      bool threadOwnsDepotLock) override {
    if (threadOwnsDepotLock) {
      this->unmarkDeletionCandidateUnderLock(asset);
    } else {
      std::lock_guard lock(this->_mutex);
      this->unmarkDeletionCandidateUnderLock(asset);
    }
  }
 
  void unmarkDeletionCandidateUnderLock(const TAssetType& asset) {
    auto it = this->_assetsByPointer.find(const_cast<TAssetType*>(&asset));
    CESIUM_ASSERT(it != this->_assetsByPointer.end());
    if (it == this->_assetsByPointer.end()) {
      return;
    }
 
    CESIUM_ASSERT(it->second != nullptr);
 
    AssetEntry& entry = *it->second;
    bool isFound = this->_deletionCandidates.contains(entry);
 
    CESIUM_ASSERT(isFound);
 
    if (isFound) {
      this->_totalDeletionCandidateMemoryUsage -= entry.sizeInDeletionList;
      this->_deletionCandidates.remove(entry);
    }
 
    // This depot is now managing at least one live asset, so keep it alive.
    this->_pKeepAlive = this;
  }
 
  struct AssetEntry
      : public CesiumUtility::ReferenceCountedThreadSafe<AssetEntry> {
    AssetEntry(TAssetKey&& key_)
        : CesiumUtility::ReferenceCountedThreadSafe<AssetEntry>(),
          key(std::move(key_)),
          pAsset(),
          maybePendingAsset(),
          errorsAndWarnings(),
          sizeInDeletionList(0),
          deletionListPointers() {}
 
    AssetEntry(const TAssetKey& key_) : AssetEntry(TAssetKey(key_)) {}
 
    TAssetKey key;
 
    std::unique_ptr<TAssetType> pAsset;
 
    std::optional<SharedFuture<CesiumUtility::ResultPointer<TAssetType>>>
        maybePendingAsset;
 
    CesiumUtility::ErrorList errorsAndWarnings;
 
    int64_t sizeInDeletionList;
 
    CesiumUtility::DoublyLinkedListPointers<AssetEntry> deletionListPointers;
 
    CesiumUtility::ResultPointer<TAssetType> toResultUnderLock() const {
      // This method is called while the calling thread already owns the depot
      // mutex. So we must take care not to lock it again, which could happen if
      // the asset is currently unreferenced and we naively create an
      // IntrusivePointer for it.
      pAsset->addReference(true);
      CesiumUtility::IntrusivePointer<TAssetType> p = pAsset.get();
      pAsset->releaseReference(true);
      return CesiumUtility::ResultPointer<TAssetType>(p, errorsAndWarnings);
    }
  };
 
  // Maps asset keys to AssetEntry instances. This collection owns the asset
  // entries.
  std::unordered_map<TAssetKey, CesiumUtility::IntrusivePointer<AssetEntry>>
      _assets;
 
  // Maps asset pointers to AssetEntry instances. The values in this map refer
  // to instances owned by the _assets map.
  std::unordered_map<TAssetType*, AssetEntry*> _assetsByPointer;
 
  // List of assets that are being considered for deletion, in the order that
  // they became unused.
  CesiumUtility::DoublyLinkedList<AssetEntry, &AssetEntry::deletionListPointers>
      _deletionCandidates;
 
  // The total amount of memory used by all assets in the _deletionCandidates
  // list.
  int64_t _totalDeletionCandidateMemoryUsage;
 
  // Mutex serializing access to _assets, _assetsByPointer, _deletionCandidates,
  // and any AssetEntry owned by this depot.
  mutable std::mutex _mutex;
 
  // The factory used to create new AssetType instances.
  std::function<FactorySignature> _factory;
 
  // This instance keeps a reference to itself whenever it is managing active
  // assets, preventing it from being destroyed even if all other references to
  // it are dropped.
  CesiumUtility::IntrusivePointer<SharedAssetDepot<TAssetType, TAssetKey>>
      _pKeepAlive;
};
 
} // namespace CesiumAsync