118 lines
3.1 KiB
C++
118 lines
3.1 KiB
C++
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
#pragma once
|
|
|
|
#include "UbaFileMapping.h"
|
|
#include "UbaLogger.h"
|
|
#include "UbaNetwork.h"
|
|
|
|
namespace uba
|
|
{
|
|
class Config;
|
|
class NetworkClient;
|
|
class NetworkServer;
|
|
class StorageImpl;
|
|
struct BinaryReader;
|
|
struct BinaryWriter;
|
|
struct ConnectionInfo;
|
|
struct MessageInfo;
|
|
struct WorkContext;
|
|
|
|
struct StorageProxyCreateInfo
|
|
{
|
|
NetworkServer& server;
|
|
NetworkClient& client;
|
|
Guid storageServerUid;
|
|
const tchar* name = TC("");
|
|
StorageImpl* localStorage = nullptr;
|
|
bool useLocalStorage = false; // Use local storage to populate proxy
|
|
|
|
void Apply(Config& config, const tchar* tableName = TC("StorageProxy"));
|
|
};
|
|
|
|
class StorageProxy
|
|
{
|
|
public:
|
|
StorageProxy(const StorageProxyCreateInfo& info);
|
|
~StorageProxy();
|
|
void PrintSummary();
|
|
|
|
u32 GetActiveFetchCount();
|
|
|
|
protected:
|
|
u16 PopId();
|
|
void PushId(u16 id);
|
|
bool HandleMessage(const ConnectionInfo& connectionInfo, const WorkContext& workContext, MessageInfo& messageInfo, BinaryReader& reader, BinaryWriter& writer);
|
|
bool HandleConnect(const ConnectionInfo& connectionInfo, BinaryReader& reader, BinaryWriter& writer);
|
|
|
|
struct BeginMessage;
|
|
bool HandleFetchBegin(const ConnectionInfo& connectionInfo, const WorkContext& workContext, MessageInfo& messageInfo, BinaryReader& reader, BinaryWriter& writer);
|
|
void HandleFetchBeginReceived(BeginMessage& m, bool error);
|
|
|
|
struct SegmentMessage;
|
|
bool HandleFetchSegment(const ConnectionInfo& connectionInfo, const WorkContext& workContext, MessageInfo& messageInfo, BinaryReader& reader, BinaryWriter& writer);
|
|
void HandleFetchSegmentReceived(SegmentMessage& mif);
|
|
|
|
bool HandleDefault(MessageInfo& messageInfo, BinaryReader& reader, BinaryWriter& writer);
|
|
|
|
struct FileEntry;
|
|
bool GetFileFromLocalStorage(u32 clientId, const CasKey& casKey, FileEntry& file, ScopedFutex& fileLock);
|
|
bool SendFetchBeginResponse(u32 clientId, FileEntry& file, BinaryWriter& writer, bool writeData = true);
|
|
bool UpdateFetch(u32 clientId, u16 fetchId, u64 segmentSize);
|
|
bool SendEnd(const CasKey& key);
|
|
|
|
static constexpr u8 ServiceId = StorageServiceId;
|
|
|
|
NetworkServer& m_server;
|
|
NetworkClient& m_client;
|
|
StorageImpl* m_localStorage;
|
|
|
|
MutableLogger m_logger;
|
|
|
|
Guid m_storageServerUid;
|
|
|
|
TString m_name;
|
|
|
|
Atomic<u32> m_inProcessClientId;
|
|
|
|
struct FileEntry
|
|
{
|
|
Futex lock;
|
|
u8* memory = nullptr;
|
|
u64 size = 0;
|
|
Atomic<u64> received;
|
|
CasKey casKey;
|
|
u32 trackId = 0;
|
|
u16 fetchId = 0;
|
|
bool storeCompressed = false;
|
|
bool sendEnd = false;
|
|
bool error = false;
|
|
bool disallowed = false;
|
|
bool available = false;
|
|
BeginMessage* beginMessage = nullptr;
|
|
Vector<SegmentMessage*> segmentMessages;
|
|
};
|
|
|
|
Futex m_filesLock;
|
|
UnorderedMap<CasKey, FileEntry> m_files;
|
|
|
|
struct ActiveFetch
|
|
{
|
|
FileEntry* file = nullptr;
|
|
u64 fetchedSize = 0;
|
|
u32 clientId = ~0u;
|
|
u32 connectionId = 0;
|
|
};
|
|
|
|
ReaderWriterLock m_activeFetchesLock;
|
|
UnorderedMap<u16, ActiveFetch> m_activeFetches;
|
|
|
|
ReaderWriterLock m_largeFileLock;
|
|
|
|
Vector<u16> m_availableIds;
|
|
u16 m_availableIdsHigh = 1;
|
|
|
|
bool m_useLocalStorage;
|
|
};
|
|
}
|