146 lines
3.9 KiB
C++
146 lines
3.9 KiB
C++
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
#pragma once
|
|
|
|
#include "UbaNetworkServer.h"
|
|
#include "UbaStorage.h"
|
|
|
|
namespace uba
|
|
{
|
|
struct StorageServerCreateInfo : StorageCreateInfo
|
|
{
|
|
StorageServerCreateInfo(NetworkServer& s, const tchar* rootDir_, LogWriter& writer) : StorageCreateInfo(rootDir_, writer), server(s) { workManager = &server; }
|
|
|
|
void Apply(const Config& config);
|
|
|
|
NetworkServer& server;
|
|
bool allowHintAsFallback = true; // Will fallback to file system to recalculate cas if cas content does not exist anymore
|
|
bool writeReceivedCasFilesToDisk = false; // Write received files to disk. This is only useful if we expect same file to be written multiple times (requires deterministic outputs ofc)
|
|
const tchar* zone = TC("");
|
|
};
|
|
|
|
class StorageServer final : public StorageImpl
|
|
{
|
|
public:
|
|
StorageServer(const StorageServerCreateInfo& info);
|
|
~StorageServer();
|
|
|
|
bool RegisterDisallowedPath(const tchar* path);
|
|
|
|
// Get the network server used by this session
|
|
NetworkServer& GetServer() { return m_server; }
|
|
|
|
using StorageImpl::StoreCasFile;
|
|
|
|
void WaitForActiveWork();
|
|
|
|
protected:
|
|
|
|
virtual bool GetZone(StringBufferBase& out) override;
|
|
virtual bool RetrieveCasFile(RetrieveResult& out, const CasKey& casKey, const tchar* hint, FileMappingBuffer* mappingBuffer = nullptr, u64 memoryMapAlignment = 1, bool allowProxy = true, u32 clientId = 0) override;
|
|
virtual bool IsDisallowedPath(const tchar* fileName) override;
|
|
virtual void SetTrace(Trace* trace, bool detailed) override;
|
|
virtual bool HasProxy(u32 clientId) override;
|
|
|
|
void OnDisconnected(u32 clientId);
|
|
bool HandleMessage(const ConnectionInfo& connectionInfo, const WorkContext& workContext, u8 messageType, BinaryReader& reader, BinaryWriter& writer);
|
|
bool WaitForWritten(CasEntry& casEntry, ScopedWriteLock& entryLock, const ConnectionInfo& connectionInfo, const tchar* hint);
|
|
bool VerifyExists(bool& outExists, CasEntry& casEntry, ScopedWriteLock& entryLock, const CasKey& casKey);
|
|
bool GetProxy(BinaryWriter& writer, u32 clientId, bool writeCasHeader);
|
|
|
|
u16 PopId();
|
|
void PushId(u16 id);
|
|
|
|
NetworkServer& m_server;
|
|
bool m_traceFetch = false;
|
|
bool m_traceStore = false;
|
|
|
|
Guid m_uid;
|
|
|
|
static constexpr u8 ServiceId = StorageServiceId;
|
|
|
|
|
|
Futex m_waitEntriesLock;
|
|
|
|
struct WaitEntry
|
|
{
|
|
WaitEntry() : done(true) {}
|
|
Event done;
|
|
bool Success = false;
|
|
u32 refCount = 0;
|
|
};
|
|
UnorderedMap<CasKey, WaitEntry> m_waitEntries;
|
|
|
|
|
|
struct ActiveStore
|
|
{
|
|
u32 clientId = ~0u;
|
|
MappedView mappedView;
|
|
FileAccessor* fileAccessor = nullptr;
|
|
CasEntry* casEntry = nullptr;
|
|
Atomic<u64> totalWritten;
|
|
Atomic<u64> recvCasTime;
|
|
u64 fileSize = 0;
|
|
u64 actualSize = 0;
|
|
};
|
|
ReaderWriterLock m_activeStoresLock;
|
|
UnorderedMap<u16, ActiveStore> m_activeStores;
|
|
|
|
|
|
struct ActiveFetch
|
|
{
|
|
u32 clientId = ~0u;
|
|
Atomic<u64> left;
|
|
CasEntry* casEntry;
|
|
Atomic<u64> sendCasTime;
|
|
|
|
FileHandle readFileHandle = InvalidFileHandle;
|
|
MappedView mappedView;
|
|
u8* memoryBegin = nullptr;
|
|
u8* memoryPos = nullptr;
|
|
bool ownsMapping = false;
|
|
|
|
void Release(StorageServer& server, const tchar* reason);
|
|
};
|
|
ReaderWriterLock m_activeFetchesLock;
|
|
UnorderedMap<u16, ActiveFetch> m_activeFetches;
|
|
Atomic<u32> m_activeUnmap;
|
|
|
|
|
|
Futex m_availableIdsLock;
|
|
Vector<u16> m_availableIds;
|
|
u16 m_availableIdsHigh = 1;
|
|
|
|
struct ProxyEntry
|
|
{
|
|
u32 clientId = ~0u;
|
|
TString zone;
|
|
TString host;
|
|
u16 port = 0;
|
|
};
|
|
ReaderWriterLock m_proxiesLock;
|
|
UnorderedMap<StringKey, ProxyEntry> m_proxies;
|
|
|
|
TString m_zone;
|
|
|
|
struct Info
|
|
{
|
|
TString zone;
|
|
TString internalAddress;
|
|
u64 storageSize = 0;
|
|
u16 proxyPort = 0;
|
|
};
|
|
ReaderWriterLock m_connectionInfoLock;
|
|
UnorderedMap<u32, Info> m_connectionInfo;
|
|
|
|
Futex m_loadCasTableLock;
|
|
|
|
Trace* m_trace = nullptr;
|
|
|
|
Vector<TString> m_disallowedPaths;
|
|
|
|
bool m_allowHintAsFallback;
|
|
bool m_writeReceivedCasFilesToDisk;
|
|
};
|
|
}
|