|
|
#pragma once
#include <iostream>
#include <atomic>
#include <memory>
#include <vector>
#include <dml/dml.hpp>
#include "util/dml-helper.hpp"
namespace dsacache { class Cache;
// cache data holds all required information on
// one cache entry and will both be stored
// internally by the cache and handed out
// as copies to the user
// this class uses its object lifetime and
// a global reference counter to allow
// thread-safe copies and resource management
class CacheData { public: using dml_handler = dml::handler<dml::mem_copy_operation, std::allocator<uint8_t>>;
private: // data source and size of the block
uint8_t* src_; size_t size_;
// global reference counting object
std::atomic<int32_t>* active_;
// global cache-location pointer
std::atomic<uint8_t*>* cache_;
// object-local incomplete cache location pointer
// which is only available in the first instance
uint8_t* incomplete_cache_;
// dml handler vector pointer which is only
// available in the first instance
std::unique_ptr<std::vector<dml_handler>> handlers_;
// deallocates the global cache-location
// and invalidates it
void Deallocate();
// checks whether there are at least two
// valid references to this object which
// is done as the cache always has one
// internally to any living instance
bool Active() const;
friend Cache; public: CacheData(uint8_t* data, const size_t size); CacheData(const CacheData& other); ~CacheData();
// waits on completion of caching operations
// for this task and is safe to be called in
// any state of the object
void WaitOnCompletion();
// returns the cache data location for this
// instance which is valid as long as the
// instance is alive - !!! this may also
// yield a nullptr !!!
uint8_t* GetDataLocation() const; }; }
inline void dsacache::CacheData::WaitOnCompletion() { // the cache data entry can be in two states
// either it is the original one which has not
// been waited for in which case the handlers
// are non-null or it is not
if (handlers_ == nullptr) { std::cout << "[-] Waiting on cache-var-update for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;
// when no handlers are attached to this cache entry we wait on a
// value change for the cache structure from nullptr to non-null
// which will either go through immediately if the cache is valid
// already or wait until the handler-owning thread notifies us
cache_->wait(nullptr);
std::cout << "[+] Finished waiting on cache-var-update for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl; } else { // when the handlers are non-null there are some DSA task handlers
// available on which we must wait here
std::cout << "[-] Waiting on handlers for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;
// abort is set if any operation encountered an error
bool abort = false;
for (auto& handler : *handlers_) { auto result = handler.get();
if (result.status != dml::status_code::ok) { std::cerr << "[x] Encountered bad status code for operation: " << dml::StatusCodeToString(result.status) << std::endl;
// if one of the copy tasks failed we abort the whole task
// after all operations are completed on it
abort = true; } }
// the handlers are cleared after all have completed
handlers_ = nullptr;
// now we act depending on whether an abort has been
// called for which signals operation incomplete
if (abort) { // store nullptr in the cache location
cache_->store(nullptr);
// then free the now incomplete cache
// TODO: it would be possible to salvage the
// TODO: operation at this point but this
// TODO: is quite complicated so we just abort
numa_free(incomplete_cache_, size_); } else { std::cout << "[+] Finished waiting on handlers for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;
// incomplete cache is now safe to use and therefore we
// swap it with the global cache state of this entry
// and notify potentially waiting threads
cache_->store(incomplete_cache_); }
// as a last step all waiting threads must
// be notified (copies of this will wait on value
// change of the cache) and the incomplete cache
// is cleared to nullptr as it is not incomplete
cache_->notify_all(); incomplete_cache_ = nullptr; } }
inline dsacache::CacheData::CacheData(uint8_t* data, const size_t size) { std::cout << "[-] New CacheData 0x" << std::hex << (uint64_t)data << std::dec << std::endl;
src_ = data; size_ = size; active_ = new std::atomic<int32_t>(1); cache_ = new std::atomic<uint8_t*>(); incomplete_cache_ = nullptr; handlers_ = std::make_unique<std::vector<dml_handler>>(); }
inline dsacache::CacheData::CacheData(const dsacache::CacheData& other) { std::cout << "[-] Copy Created for CacheData 0x" << std::hex << (uint64_t)other.src_ << std::dec << std::endl;
// we copy the ptr to the global atomic reference counter
// and increase the amount of active references
active_ = other.active_; const int current_active = active_->fetch_add(1);
// source and size will be copied too
// as well as the reference to the global
// atomic cache pointer
src_ = other.src_; size_ = other.size_; cache_ = other.cache_;
// incomplete cache and handlers will not
// be copied because only the first instance
// will wait on the completion of handlers
incomplete_cache_ = nullptr; handlers_ = nullptr; }
inline dsacache::CacheData::~CacheData() { std::cout << "[-] Destructor for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;
// if this is the first instance of this cache structure
// and it has not been waited on and is now being destroyed
// we must wait on completion here to ensure the cache
// remains in a valid state
if (handlers_ != nullptr) { WaitOnCompletion(); }
// due to fetch_sub returning the preivously held value
// we must subtract one locally to get the current value
const int32_t v = active_->fetch_sub(1) - 1;
// if the returned value is zero or lower
// then we must execute proper deletion
// as this was the last reference
if (v <= 0) { std::cout << "[!] Full Destructor for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;
Deallocate();
delete active_; delete cache_; } }
inline void dsacache::CacheData::Deallocate() { std::cout << "[!] Deallocating for CacheData 0x" << std::hex << (uint64_t)src_ << std::dec << std::endl;
// although deallocate should only be called from
// a safe context to do so, it can not hurt to
// defensively perform the operation atomically
uint8_t* cache_local = cache_->exchange(nullptr); if (cache_local != nullptr) numa_free(cache_local, size_); }
inline uint8_t* dsacache::CacheData::GetDataLocation() const { return cache_->load(); }
inline bool dsacache::CacheData::Active() const { // this entry is active if more than one
// reference exists to it, as the Cache
// will always keep one internally until
// the entry is cleared from cache
return active_->load() > 1; }
|