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 {

struct SharedAssetContext {
  AsyncSystem asyncSystem;

  std::shared_ptr<IAssetAccessor> pAssetAccessor;
};

template <
    typename TAssetType,
    typename TAssetKey,
    typename TContext = SharedAssetContext>
class CESIUMASYNC_API SharedAssetDepot
    : public CesiumUtility::ReferenceCountedThreadSafe<
          SharedAssetDepot<TAssetType, TAssetKey, TContext>>,
      public CesiumUtility::IDepotOwningAsset<TAssetType> {
public:
  std::atomic<int64_t> inactiveAssetSizeLimitBytes =
      static_cast<int64_t>(16 * 1024 * 1024);

  using FactorySignature =
      CesiumAsync::Future<CesiumUtility::ResultPointer<TAssetType>>(
          const TContext& context,
          const TAssetKey& key);

  SharedAssetDepot(std::function<FactorySignature> factory);
 
  virtual ~SharedAssetDepot();

  SharedFuture<CesiumUtility::ResultPointer<TAssetType>>
  getOrCreate(const TContext& context, const TAssetKey& assetKey);

  bool invalidate(const TAssetKey& assetKey);

  bool invalidate(TAssetType& asset);

  size_t getAssetCount() const;

  size_t getActiveAssetCount() const;

  size_t getInactiveAssetCount() const;

  int64_t getInactiveAssetTotalSizeBytes() const;
 
  // Disable copy
  void operator=(
      const SharedAssetDepot<TAssetType, TAssetKey, TContext>& other) = delete;
 
private:
  struct LockHolder;

  LockHolder lock() const;

  void markDeletionCandidate(const TAssetType& asset, bool threadOwnsDepotLock)
      override;
 
  void markDeletionCandidateUnderLock(const TAssetType& asset);

  void unmarkDeletionCandidate(
      const TAssetType& asset,
      bool threadOwnsDepotLock) override;
 
  void unmarkDeletionCandidateUnderLock(const TAssetType& asset);

  bool invalidateUnderLock(LockHolder&& lock, const TAssetKey& assetKey);

  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;
  };
 
  // Manages the depot's mutex. Also ensures, via IntrusivePointer, that the
  // depot won't be destroyed while the lock is held.
  struct LockHolder {
    LockHolder(
        const CesiumUtility::IntrusivePointer<const SharedAssetDepot>& pDepot);
    ~LockHolder();
    void unlock();
 
  private:
    // These two fields _must_ be declared in this order to guarantee that the
    // mutex is released before the depot pointer. Releasing the depot pointer
    // could destroy the depot, and that will be disastrous if the lock is still
    // held.
    CesiumUtility::IntrusivePointer<const SharedAssetDepot> pDepot;
    std::unique_lock<std::mutex> lock;
  };
 
  // 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;
 
  // The number of assets that have been invalidated but that have not been
  // deleted yet. Such assets hold a pointer to the depot, so the depot must be
  // kept alive for their entire lifetime.
  int64_t _liveInvalidatedAssets;
 
  // 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, TContext>>
      _pKeepAlive;
};
 
template <typename TAssetType, typename TAssetKey, typename TContext>
SharedAssetDepot<TAssetType, TAssetKey, TContext>::SharedAssetDepot(
    std::function<FactorySignature> factory)
    : _assets(),
      _assetsByPointer(),
      _deletionCandidates(),
      _totalDeletionCandidateMemoryUsage(0),
      _liveInvalidatedAssets(0),
      _mutex(),
      _factory(std::move(factory)),
      _pKeepAlive(nullptr) {}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
SharedAssetDepot<TAssetType, TAssetKey, TContext>::~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->_liveInvalidatedAssets == 0);
  CESIUM_ASSERT(this->_assets.size() == this->_deletionCandidates.size());
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
SharedFuture<CesiumUtility::ResultPointer<TAssetType>>
SharedAssetDepot<TAssetType, TAssetKey, TContext>::getOrCreate(
    const TContext& context,
    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.
  LockHolder lock = this->lock();
 
  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 context.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 = context.asyncSystem.template createPromise<void>();
 
  // We haven't loaded or started to load this asset yet.
  // Let's do that now.
  CesiumUtility::IntrusivePointer<
      SharedAssetDepot<TAssetType, TAssetKey, TContext>>
      pDepot = this;
  CesiumUtility::IntrusivePointer<AssetEntry> pEntry = new AssetEntry(assetKey);
 
  auto future =
      promise.getFuture()
          .thenImmediately([pDepot, pEntry, context]() {
            return pDepot->_factory(context, 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) {
                LockHolder lock = pDepot->lock();
 
                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;
 
  [[maybe_unused]] bool added = this->_assets.emplace(assetKey, pEntry).second;
 
  // 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;
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
bool SharedAssetDepot<TAssetType, TAssetKey, TContext>::invalidate(
    const TAssetKey& assetKey) {
  LockHolder lock = this->lock();
  return this->invalidateUnderLock(std::move(lock), assetKey);
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
bool SharedAssetDepot<TAssetType, TAssetKey, TContext>::invalidate(
    TAssetType& asset) {
  LockHolder lock = this->lock();
 
  auto it = this->_assetsByPointer.find(&asset);
  if (it == this->_assetsByPointer.end())
    return false;
 
  AssetEntry* pEntry = it->second;
  CESIUM_ASSERT(pEntry);
 
  return this->invalidateUnderLock(std::move(lock), pEntry->key);
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
size_t
SharedAssetDepot<TAssetType, TAssetKey, TContext>::getAssetCount() const {
  LockHolder lock = this->lock();
  return this->_assets.size();
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
size_t
SharedAssetDepot<TAssetType, TAssetKey, TContext>::getActiveAssetCount() const {
  LockHolder lock = this->lock();
  return this->_assets.size() - this->_deletionCandidates.size();
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
size_t
SharedAssetDepot<TAssetType, TAssetKey, TContext>::getInactiveAssetCount()
    const {
  LockHolder lock = this->lock();
  return this->_deletionCandidates.size();
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
int64_t SharedAssetDepot<TAssetType, TAssetKey, TContext>::
    getInactiveAssetTotalSizeBytes() const {
  LockHolder lock = this->lock();
  return this->_totalDeletionCandidateMemoryUsage;
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
typename SharedAssetDepot<TAssetType, TAssetKey, TContext>::LockHolder
SharedAssetDepot<TAssetType, TAssetKey, TContext>::lock() const {
  return LockHolder{this};
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
void SharedAssetDepot<TAssetType, TAssetKey, TContext>::markDeletionCandidate(
    const TAssetType& asset,
    bool threadOwnsDepotLock) {
  if (threadOwnsDepotLock) {
    this->markDeletionCandidateUnderLock(asset);
  } else {
    LockHolder lock = this->lock();
    this->markDeletionCandidateUnderLock(asset);
  }
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
void SharedAssetDepot<TAssetType, TAssetKey, TContext>::
    markDeletionCandidateUnderLock(const TAssetType& asset) {
  if (asset._isInvalidated) {
    // This asset is no longer tracked by the depot, so delete it.
    --this->_liveInvalidatedAssets;
    delete &asset;
 
    // 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->_liveInvalidatedAssets == 0) {
      this->_pKeepAlive.reset();
    }
 
    return;
  }
 
  // Verify that the reference count is still zero.
  // See: https://github.com/CesiumGS/cesium-native/issues/1073
  if (asset._referenceCount != 0) {
    return;
  }
 
  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->_liveInvalidatedAssets == 0) {
    this->_pKeepAlive.reset();
  }
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
void SharedAssetDepot<TAssetType, TAssetKey, TContext>::unmarkDeletionCandidate(
    const TAssetType& asset,
    bool threadOwnsDepotLock) {
  if (threadOwnsDepotLock) {
    this->unmarkDeletionCandidateUnderLock(asset);
  } else {
    LockHolder lock = this->lock();
    this->unmarkDeletionCandidateUnderLock(asset);
  }
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
void SharedAssetDepot<TAssetType, TAssetKey, TContext>::
    unmarkDeletionCandidateUnderLock(const TAssetType& asset) {
  // This asset better not already be invalidated. That would imply this asset
  // was resurrected after its reference count hit zero. This should only be
  // possible if the asset depot returned a pointer to the asset, which it
  // will not do for one that is invalidated.
  CESIUM_ASSERT(!asset._isInvalidated);
 
  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);
 
  // The asset won't necessarily be found in the deletionCandidates set.
  // See: https://github.com/CesiumGS/cesium-native/issues/1073
  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;
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
bool SharedAssetDepot<TAssetType, TAssetKey, TContext>::invalidateUnderLock(
    LockHolder&& lock,
    const TAssetKey& assetKey) {
  auto it = this->_assets.find(assetKey);
  if (it == this->_assets.end())
    return false;
 
  AssetEntry* pEntry = it->second.get();
  CESIUM_ASSERT(pEntry);
 
  // This will remove the asset from the deletion candidates list, if it's
  // there.
  CesiumUtility::ResultPointer<TAssetType> assetResult =
      pEntry->toResultUnderLock();
 
  bool wasInvalidated = false;
 
  if (assetResult.pValue) {
    if (!assetResult.pValue->_isInvalidated) {
      wasInvalidated = true;
      assetResult.pValue->_isInvalidated = true;
      ++this->_liveInvalidatedAssets;
    }
    this->_assetsByPointer.erase(assetResult.pValue.get());
  }
 
  // Detach the asset from the AssetEntry, so that its lifetime is controlled by
  // reference counting.
  pEntry->pAsset.release();
 
  // Remove the asset entry. This won't immediately delete the asset, because
  // `assetResult` above still holds a reference to it. But once that goes out
  // of scope, too, the asset _may_ be destroyed.
  this->_assets.erase(it);
 
  // Unlock the mutex before allowing `assetResult` to go out of scope. When it
  // goes out of scope, the asset may be destroyed. If it is, that would cause
  // us to try to re-enter the lock, which is not allowed.
  lock.unlock();
 
  return wasInvalidated;
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
CesiumUtility::ResultPointer<TAssetType>
SharedAssetDepot<TAssetType, TAssetKey, TContext>::AssetEntry::
    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.
  CesiumUtility::IntrusivePointer<TAssetType> p = nullptr;
  if (pAsset) {
    pAsset->addReference(true);
    p = pAsset.get();
    pAsset->releaseReference(true);
  }
  return CesiumUtility::ResultPointer<TAssetType>(p, errorsAndWarnings);
}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
SharedAssetDepot<TAssetType, TAssetKey, TContext>::LockHolder::LockHolder(
    const CesiumUtility::IntrusivePointer<const SharedAssetDepot>& pDepot_)
    : pDepot(pDepot_), lock(pDepot_->_mutex) {}
 
template <typename TAssetType, typename TAssetKey, typename TContext>
SharedAssetDepot<TAssetType, TAssetKey, TContext>::LockHolder::~LockHolder() =
    default;
 
template <typename TAssetType, typename TAssetKey, typename TContext>
void SharedAssetDepot<TAssetType, TAssetKey, TContext>::LockHolder::unlock() {
  this->lock.unlock();
}
 
} // namespace CesiumAsync