DataFetcherManager.cpp

#include <cassert>
#include "zstd.h"
#include <bit>
#include <shared_mutex>
#include <brotli/decode.h>
 
#include "Foundation/Logging/Logger.h"
#include "Foundation/Platform/MemoryBufferBackedFileContents.h"
#include "Context.h"
#include "Services/TaskManager.h"
 
#include "DataFetcherManager.h"
 
#include "DataFetcherBase.h"
 
namespace {
  struct Item {
    Cogs::Core::DataFetcherBase* fetcher = nullptr;
    int priority = 0;
    uint32_t id = 0;
  };
 
  std::shared_mutex fetcherLock;
  std::vector<Item> fetchers;
  Cogs::Logging::Log logger = Cogs::Logging::getLogger("DataFetcherManager");
 
  uint32_t currentFetcherId = 0;
  uint32_t currentRequestId = 0;
  constexpr uint32_t fetcherIdBits = 7;
  constexpr uint32_t fetcherIdMask = (1u << fetcherIdBits) - 1;
}
 
void Cogs::Core::DataFetcherManager::addDataFetcher(DataFetcherBase* dataFetcher, int priority)
{
  std::unique_lock guard(fetcherLock);  // Needs exclusive access as it modifies the list
  assert((fetchers.size() < 100) && "Manage the fetchers smarter");
  for (const auto& fetcher : fetchers) {
    if (fetcher.fetcher == dataFetcher) {
      LOG_ERROR(logger, "addDataFetcher - ignoring duplicate DataFetcher");
      return;
    }
  }
 
  fetchers.push_back({});
  // Insertion sort, no point in full blown introsort for just a few elements.
  size_t i = fetchers.size() - 1;
  while (0 < i && fetchers[i - 1].priority < priority) {
    fetchers[i] = fetchers[i - 1];
    i--;
  }
  fetchers[i].fetcher = dataFetcher;
  fetchers[i].priority = priority;
 
  if (currentFetcherId == (NoFetchId & fetcherIdMask)) { currentFetcherId = (currentFetcherId + 1) & fetcherIdMask; }
  fetchers[i].id = currentFetcherId;
  currentFetcherId = (currentFetcherId + 1) & fetcherIdMask;
}
 
void Cogs::Core::DataFetcherManager::removeDataFetcher(DataFetcherBase* dataFetcher) {
  std::unique_lock guard(fetcherLock);  // Needs exclusive access as it modifies the list
  for (auto i = fetchers.begin(), e = fetchers.end(); i != e; ++i) {
    if (i->fetcher == dataFetcher) {
      fetchers.erase(i);
      return;
    }
  }
 
  LOG_ERROR(logger, "removeDataFetcher - unknown DataFetcher");
}
 
namespace {
  using namespace Cogs::Core;
 
  std::unique_ptr<Cogs::FileContents> brotliDecompress(std::unique_ptr<Cogs::FileContents>&& compressed)
  {
    BrotliDecoderState* decoder = BrotliDecoderCreateInstance(nullptr, nullptr, nullptr); // No custom memory allocator
    assert(decoder);
 
    const uint8_t* nextIn = compressed->ptr;
    size_t availableIn = compressed->size;
    Cogs::Memory::MemoryBuffer decompressed;
    while (!BrotliDecoderIsFinished(decoder)) {
      size_t availableOut = 0;
      BrotliDecoderResult rv = BrotliDecoderDecompressStream(decoder,
                                                             &availableIn,
                                                             &nextIn,
                                                             &availableOut,
                                                             nullptr,                 // next output buffer
                                                             nullptr);                // output so far
 
      if (rv == BROTLI_DECODER_RESULT_ERROR) {
        LOG_ERROR(logger, "Parsing bin: BrotliDecoderDecompressStream failed: %s", BrotliDecoderErrorString(static_cast<BrotliDecoderErrorCode>(rv)));
        BrotliDecoderDestroyInstance(decoder);
        return std::unique_ptr<Cogs::FileContents>();
      }
 
      size_t outputSize = 0;
      const uint8_t* outputBuffer = BrotliDecoderTakeOutput(decoder, &outputSize);
      bool success = decompressed.write(outputBuffer, outputSize);
      assert(success);
    }
    BrotliDecoderDestroyInstance(decoder);
 
    LOG_TRACE(logger, "Brotli-decompressed %zu -> %zu bytes in transit", compressed->size, decompressed.size());
 
    return std::make_unique<Cogs::MemoryBufferBackedFileContents>(std::move(decompressed),
                                                                  compressed->origin().to_string(),
                                                                  compressed->hints & ~Cogs::FileContentsHints::BrotliDecompress);
  }
 
  std::unique_ptr<Cogs::FileContents> zstdDecompress(std::unique_ptr<Cogs::FileContents>&& compressed)
  {
    ZSTD_DCtx* zstdctx = ZSTD_createDCtx();
 
    unsigned long long outSize = ZSTD_getFrameContentSize(compressed->ptr, compressed->size);
 
    if (outSize == ZSTD_CONTENTSIZE_UNKNOWN) {
      LOG_ERROR(logger, "Parsing bin: ZSTD_GetFrameContentSize returned ZSTD_CONTENTSIZE_UNKNOWN");
      ZSTD_freeDCtx(zstdctx);
      return std::unique_ptr<Cogs::FileContents>();
    }
    else if (outSize == ZSTD_CONTENTSIZE_ERROR) {
      LOG_ERROR(logger, "Parsing bin: ZSTD_GetFrameContentSize returned ZSTD_CONTENTSIZE_ERROR");
      ZSTD_freeDCtx(zstdctx);
      return std::unique_ptr<Cogs::FileContents>();
    }
 
    Cogs::Memory::MemoryBuffer decompressed;
    decompressed.resize(outSize);
    size_t status = ZSTD_decompressDCtx(zstdctx,
                                        decompressed.data(), decompressed.size(),
                                        compressed->ptr, compressed->size);
    ZSTD_freeDCtx(zstdctx);
 
    if (ZSTD_isError(status)) {
      LOG_ERROR(logger, "Parsing bin: ZSTD_decompressDCtx  failed: %s", ZSTD_getErrorName(status));
      return std::unique_ptr<Cogs::FileContents>();
    }
 
    LOG_TRACE(logger, "ZStd-decompressed %zu -> %zu bytes in transit", compressed->size, decompressed.size());
 
    return std::make_unique<Cogs::MemoryBufferBackedFileContents>(std::move(decompressed),
                                                                  compressed->origin().to_string(),
                                                                  compressed->hints & ~Cogs::FileContentsHints::ZStdDecompress);
  }
 
  void maybeDecompressAndInvokeHandler(const Cogs::FileContents::Callback& callback, std::unique_ptr<Cogs::FileContents> && data)
  {
    if (!data) {
      // No data was loaded.  Just pass this along.
      callback(std::move(data));
      return;
    }
 
    // Check if the data should be decompressed in transit
    if ((data->hints & Cogs::FileContentsHints::BrotliDecompress) == Cogs::FileContentsHints::BrotliDecompress) {
      callback(brotliDecompress(std::move(data)));
      return;
    }
 
    else if ((data->hints & Cogs::FileContentsHints::ZStdDecompress) == Cogs::FileContentsHints::ZStdDecompress) {
 
      // Since ZStd compressed files can be identified via a magic number at the beginning, we check
      // for that before decompressing it. Hence we can add the FileContentsHints::ZStdDecompress
      // speculatively and transparently just compress if it is actually compressed, something that
      // is used for compressed asset hierarchies (see notes in Documentation/AssetSystem.md)
      static_assert(std::endian::native == std::endian::little); // TODO: Handle case below in the unlikely event that cogs is compiled on a non little-endian architecture.
      if (4 <= data->size && *reinterpret_cast<const uint32_t*>(data->ptr) == 0xFD2FB528) {
        callback(zstdDecompress(std::move(data)));
        return;
      }
    }
 
    callback(std::move(data));
  }
 
}
 
Cogs::Core::DataFetcherManager::FetchId
Cogs::Core::DataFetcherManager::fetchAsync(Context* context, const std::string& fileName, const FileContents::Callback& callback, uint64_t offset, uint64_t size, bool cancellable, FileContentsHints hints)
{
  assert(context);
 
  // Wrap the callback in our own callback that marshalls to the resource queue if we are on
  // the main thread and tries to decompress the data if it is known to be compressed.
  FileContents::Callback moveToResourceQueueHandler =
    [context, callback](std::unique_ptr<Cogs::FileContents>&& data)
    {
      TaskManager& tm = *context->taskManager;
      if (tm.onMainThread() && tm.getQueueConcurrency(TaskManager::ResourceQueue)) {
        tm.enqueue(TaskManager::ResourceQueue,
                  [callback, dataPtr = data.release()]()
                  {
                    maybeDecompressAndInvokeHandler(callback, std::unique_ptr<Cogs::FileContents>(dataPtr));
                  });
      }
      else {
        maybeDecompressAndInvokeHandler(callback, std::move(data));
      }
    };
 
 
  // Try to pass this on to the fetchers
  {
    std::shared_lock guard(fetcherLock);  // Can use shared access since it only reads
    for (Item& item : fetchers) {
 
      FetchId cancellationId = NoFetchId;
      if (cancellable) {
        cancellationId = (currentRequestId << fetcherIdBits) | item.id;
        assert(cancellationId != NoFetchId); // id never matches the lower NoFetchId bit pattern
      }
 
      if (item.fetcher->fetchAsync(fileName, moveToResourceQueueHandler, offset, size, cancellationId, hints)) {
        if (cancellable) {
          currentRequestId++;
        }
        return cancellationId;
      }
    }
  }
 
#ifdef EMSCRIPTEN
  IO::readFileAsync(fileName, moveToResourceQueueHandler, offset, size, hints);
#else
  // On desktop, we do syncronous file reads in a resource queue worker.
  // Note that we pass the original callback since we already are in a resource worker
  context->taskManager->enqueue(TaskManager::ResourceQueue,
                                [fileName, callback, hints, offset, size]()// mutable
                                {
                                  maybeDecompressAndInvokeHandler(callback, IO::readFileSync(fileName, offset, size, hints));
                                });
 
#endif
  return NoFetchId;
}
 
Cogs::FileContents::Ptr Cogs::Core::DataFetcherManager::fetchSync(const std::string& fileName, uint64_t offset, uint64_t size, FileContentsHints hints)
{
  for (Item& item : fetchers) {
    FileContents::Ptr contents = item.fetcher->fetchSync(fileName, offset, size, hints);
 
    if (contents) {
      return contents;
    }
  }
  return IO::readFileSync(fileName, offset, size, hints);
}
 
void Cogs::Core::DataFetcherManager::cancelAsyncFetch(Context* /*context*/, FetchId id)
{
  if (id == NoFetchId) return;
 
  for (Item& item : fetchers) {
    if (item.id == (id & fetcherIdMask)) {
      item.fetcher->cancelAsyncFetch(id);
      return;
    }
  }
}