Files
UnrealEngine/Engine/Source/Developer/Horde/Private/Compute/ComputeBuffer.cpp
2025-05-18 13:04:45 +08:00

1208 lines
28 KiB
C++

// Copyright Epic Games, Inc. All Rights Reserved.
#include "Compute/ComputeBuffer.h"
#include "Compute/ComputePlatform.h"
#include "../HordePlatform.h"
#include <assert.h>
#include <iostream>
#include <assert.h>
#include <stdio.h>
#include <string.h>
struct FComputeBufferDetail
{
static const int HeaderSize = (2 + FComputeBuffer::MaxReaders + FComputeBuffer::MaxChunks) * sizeof(long long);
enum class EWriteState : unsigned long long
{
// Writer has moved to the next chunk
MovedToNext = 0,
// Chunk is still being appended to
Writing = 2,
// This chunk marks the end of the stream
Complete = 3,
};
union FChunkState
{
public:
long long Value;
FChunkState()
: Value(0)
{ }
FChunkState(long long InValue)
: Value(InValue)
{ }
// Constructor
FChunkState(EWriteState WriteState, int ReaderFlags, int Length)
: Value((unsigned long long)Length | ((unsigned long long)ReaderFlags << 31) | ((unsigned long long)WriteState << 62))
{ }
// Written length of this chunk
int GetLength() const { return (int)(Value & 0x7fffffff); }
// Set of flags which are set for each reader that still has to read from a chunk
int GetReaderFlags() const { return (int)((Value >> 31) & 0x7fffffff); }
// State of the writer
EWriteState GetWriteState() const { return (EWriteState)((unsigned long long)Value >> 62); }
// Test whether a particular reader is still referencing the chunk
bool HasReaderFlag(int ReaderIdx) const { return (Value & (1ULL << (31 + ReaderIdx))) != 0; }
private:
struct // For debugging purposes only; non-portable assumption of bitfield layout
{
unsigned long long Length : 31;
unsigned long long ReaderFlags : 31;
EWriteState WriteState : 2;
} Fields;
};
struct FChunkStatePtr
{
public:
FChunkStatePtr(volatile FChunkState* InChunkState)
: State(InChunkState)
{ }
// Read the state value from memory
FChunkState Get() const
{
return FChunkState(FPlatformAtomics::AtomicRead(&State->Value));
}
// Read the state value from memory
void Set(FChunkState NewState)
{
FPlatformAtomics::AtomicStore(&State->Value, NewState.Value);
}
// Attempt to update the chunk state
bool TryUpdate(FChunkState PrevState, FChunkState NextState)
{
return FPlatformAtomics::InterlockedCompareExchange(&State->Value, NextState.Value, PrevState.Value) == PrevState.Value;
}
// Append data to the chunk
void Append(long long Length)
{
FPlatformAtomics::InterlockedAdd(&State->Value, Length);
}
// Mark this chunk as the end of the stream
void MarkComplete()
{
FPlatformAtomics::InterlockedOr(&State->Value, FChunkState(EWriteState::Complete, 0, 0).Value);
}
// Start reading the chunk with the given reader
void StartReading(int ReaderIdx)
{
FPlatformAtomics::InterlockedOr(&State->Value, FChunkState((EWriteState)0, 1 << ReaderIdx, 0).Value);
}
// Clear the reader flag
void FinishReading(int ReaderIdx)
{
FPlatformAtomics::InterlockedAnd(&State->Value, ~FChunkState((EWriteState)0, 1 << ReaderIdx, 0).Value);
}
// Move to the next chunk
void FinishWriting()
{
FPlatformAtomics::InterlockedAnd(&State->Value, ~FChunkState(EWriteState::Writing, 0, 0).Value);
}
private:
volatile FChunkState* const State;
};
union FReaderState
{
long long Value;
FReaderState()
: Value(0)
{ }
FReaderState(long long InValue)
: Value(InValue)
{ }
FReaderState(int ChunkIdx, int Offset, int RefCount, bool Detached)
: Value((long long)(unsigned int)Offset | ((long long)(unsigned int)ChunkIdx << 32) | ((long long)(unsigned int)RefCount << 40) | (long long)(Detached ? (1ULL << 63) : 0))
{ }
int GetOffset() const
{
return (int)(Value & 0xffffffff);
}
int GetChunkIdx() const
{
return (int)((Value >> 32) & 0xff);
}
int GetRefCount() const
{
return (int)((Value >> 40) & 0x7fff);
}
bool IsDetached() const
{
return (Value & (1ULL << 63)) != 0;
}
private:
struct // For debugging purposes only; non-portable assumption of bitfield layout
{
unsigned long long Offset : 32;
unsigned long long ChunkIdx : 8;
unsigned long long RefCount : 23;
unsigned long long Detached : 1;
} Fields;
};
struct FReaderStatePtr
{
public:
FReaderStatePtr(volatile FReaderState* InReaderState)
: State(InReaderState)
{ }
// Read the state value from memory
FReaderState Get() const
{
return FReaderState(FPlatformAtomics::AtomicRead(&State->Value));
}
// Read the state value from memory
void Set(FReaderState NewState)
{
FPlatformAtomics::AtomicStore(&State->Value, NewState.Value);
}
// Attempt to update the chunk state
bool TryUpdate(FReaderState PrevState, FReaderState NextState)
{
return FPlatformAtomics::InterlockedCompareExchange(&State->Value, NextState.Value, PrevState.Value) == PrevState.Value;
}
private:
volatile FReaderState* const State;
};
union FWriterState
{
long long Value;
FWriterState()
: Value(0)
{ }
FWriterState(long long InValue)
: Value(InValue)
{ }
FWriterState(int ChunkIdx, int ReaderFlags, int RefCount, bool HasWrapped)
: Value((long long)(unsigned int)ChunkIdx | ((long long)(unsigned int)ReaderFlags << 32) | ((long long)(unsigned int)RefCount << 48) | (HasWrapped ? (1ULL << 63) : 0))
{ }
int GetChunkIdx() const { return Value & 0x7fffffff; }
int GetReaderFlags() const { return ((unsigned long long)Value >> 32) & 0xffff; }
int GetRefCount() const { return ((unsigned long long)Value >> 48) & 0x7fff; }
bool HasWrapped() const { return (Value & (1ULL << 63)) != 0; }
private:
struct // For debugging purposes only; non-portable assumption of bitfield layout
{
int ChunkIdx : 32;
int ReaderFlags : 8;
int RefCount : 23;
unsigned int Wrapped : 1;
} Fields;
};
struct FWriterStatePtr
{
public:
FWriterStatePtr(volatile FWriterState* InState)
: State(InState)
{ }
FWriterState Get() const
{
return FWriterState(FPlatformAtomics::AtomicRead(&State->Value));
}
void Set(FWriterState NewState)
{
FPlatformAtomics::AtomicStore(&State->Value, NewState.Value);
}
bool TryUpdate(FWriterState PrevValue, FWriterState NextValue)
{
return FPlatformAtomics::InterlockedCompareExchange(&State->Value, NextValue.Value, PrevValue.Value) == PrevValue.Value;
}
private:
volatile FWriterState* State;
};
struct FHeader
{
public:
const unsigned short int NumReaders;
const unsigned short int NumChunks;
const int ChunkLength;
FHeader(int InNumReaders, int InNumChunks, int InChunkLength)
: NumReaders(InNumReaders)
, NumChunks(InNumChunks)
, ChunkLength(InChunkLength)
{
static_assert(sizeof(FWriterState) == sizeof(long long), "Incorrect size of FWriterState");
static_assert(sizeof(FReaderState) == sizeof(long long), "Incorrect size of FReaderState");
static_assert(sizeof(FChunkState) == sizeof(long long), "Incorrect size of FChunkState");
static_assert(sizeof(FHeader) == HeaderSize, "Incorrect size of FHeader");
GetChunkStatePtr(0).Set(FChunkState(EWriteState::Writing, 0, 0));
}
FWriterStatePtr GetWriterStatePtr()
{
return FWriterStatePtr(&Writer);
}
FReaderStatePtr GetReaderStatePtr(int ReaderIdx)
{
return FReaderStatePtr(&Readers[ReaderIdx]);
}
FChunkStatePtr GetChunkStatePtr(int ChunkIdx)
{
return FChunkStatePtr(&Chunks[ChunkIdx]);
}
private:
FWriterState Writer;
FReaderState Readers[FComputeBuffer::MaxReaders];
FChunkState Chunks[FComputeBuffer::MaxChunks];
};
char Name[FComputeBuffer::MaxNameLength];
FHeader* Header;
unsigned char* ChunkPtrs[FComputeBuffer::MaxChunks];
FComputeEvent WriterEvent;
FComputeEvent ReaderEvents[FComputeBuffer::MaxReaders];
int32 RefCount;
FComputeBufferDetail(const char* InName)
: Header(nullptr)
, ChunkPtrs{ nullptr, }
, RefCount(1)
{
static_assert(sizeof(FChunkStatePtr) == sizeof(long long), "Incorrect size of FChunkStatePtr; check union is declared correctly.");
static_assert(sizeof(FWriterStatePtr) == sizeof(long long), "Incorrect size of FWriterStatePtr; check union is declared correctly.");
FCStringAnsi::Strncpy(Name, InName, FComputeBuffer::MaxNameLength);
}
~FComputeBufferDetail()
{
Header = nullptr;
}
void AddRef()
{
FPlatformAtomics::InterlockedIncrement(&RefCount);
}
void Release()
{
if (FPlatformAtomics::InterlockedDecrement(&RefCount) == 0)
{
delete this;
}
}
static TUniquePtr<FComputeBufferDetail> CreateNew(const FComputeBuffer::FParams& Params)
{
long long Capacity = sizeof(FHeader) + (Params.NumChunks * sizeof(unsigned int)) + (Params.NumChunks * Params.ChunkLength);
const char* Name = Params.Name;
char BaseNameBuffer[FComputeBuffer::MaxNameLength];
if (Name == nullptr)
{
FHordePlatform::CreateUniqueName(BaseNameBuffer, FComputeBuffer::MaxNameLength);
Name = BaseNameBuffer;
}
char NameBuffer[FComputeBuffer::MaxNameLength];
snprintf(NameBuffer, FComputeBuffer::MaxNameLength, "%s_M", Name);
TUniquePtr<FComputeBufferDetail> Detail = MakeUnique<FComputeBufferDetail>(Name);
if (!Detail->MemoryMappedFile.Create(Params.Anonymous ? nullptr : NameBuffer, Capacity))
{
checkf(false, TEXT("Failed to create memory mapped file %hs"), NameBuffer);
return nullptr;
}
void* Pointer = Detail->MemoryMappedFile.GetPointer();
if (Pointer == nullptr)
{
checkf(false, TEXT("Failed to get pointer to memory mapping %hs"), NameBuffer);
return nullptr;
}
Detail->Header = new(Pointer) FHeader(Params.NumReaders, Params.NumChunks, Params.ChunkLength);
snprintf(NameBuffer, FComputeBuffer::MaxNameLength, "%s_W", Name);
if (!Detail->WriterEvent.Create(Params.Anonymous ? nullptr : NameBuffer))
{
checkf(false, TEXT("Failed to create writer event %hs in memory mapping"), NameBuffer);
return nullptr;
}
for (int ReaderIdx = 0; ReaderIdx < Detail->Header->NumReaders; ReaderIdx++)
{
snprintf(NameBuffer, FComputeBuffer::MaxNameLength, "%s_R%d", Name, ReaderIdx);
if (!Detail->ReaderEvents[ReaderIdx].Create(Params.Anonymous ? nullptr : NameBuffer))
{
checkf(false, TEXT("Failed to create reader event %hs in memory mapping"), NameBuffer);
return nullptr;
}
}
InitShared(Detail.Get());
return Detail;
}
static TUniquePtr<FComputeBufferDetail> OpenExisting(const char* Name)
{
char NameBuffer[FComputeBuffer::MaxNameLength];
snprintf(NameBuffer, FComputeBuffer::MaxNameLength, "%s_M", Name);
TUniquePtr<FComputeBufferDetail> Detail = MakeUnique<FComputeBufferDetail>(Name);
if (!Detail->MemoryMappedFile.OpenExisting(NameBuffer))
{
return nullptr;
}
Detail->Header = (FHeader*)Detail->MemoryMappedFile.GetPointer();
if (Detail->Header == nullptr)
{
return nullptr;
}
snprintf(NameBuffer, FComputeBuffer::MaxNameLength, "%s_W", Name);
if (!Detail->WriterEvent.OpenExisting(NameBuffer))
{
return nullptr;
}
for (int ReaderIdx = 0; ReaderIdx < Detail->Header->NumReaders; ReaderIdx++)
{
snprintf(NameBuffer, FComputeBuffer::MaxNameLength, "%s_R%d", Name, ReaderIdx);
if (!Detail->ReaderEvents[ReaderIdx].OpenExisting(NameBuffer))
{
return nullptr;
}
}
InitShared(Detail.Get());
return Detail;
}
int CreateReader()
{
for(int ReaderIdx = 0; ReaderIdx < Header->NumReaders; ReaderIdx++)
{
FReaderStatePtr ReaderStatePtr = Header->GetReaderStatePtr(ReaderIdx);
for (;;)
{
FReaderState ReaderState = ReaderStatePtr.Get();
if (ReaderState.GetRefCount() > 0)
{
break;
}
if (ReaderStatePtr.TryUpdate(ReaderState, FReaderState(0, 0, 1, false)))
{
FWriterStatePtr WriterStatePtr = Header->GetWriterStatePtr();
for (;;)
{
FWriterState WriterState = WriterStatePtr.Get();
check(!WriterState.HasWrapped());
if (WriterStatePtr.TryUpdate(WriterState, FWriterState(WriterState.GetChunkIdx(), WriterState.GetReaderFlags() | (1 << ReaderIdx), WriterState.GetRefCount(), WriterState.HasWrapped())))
{
for (int WriteChunkIdx = 0; WriteChunkIdx <= WriterState.GetChunkIdx(); WriteChunkIdx++)
{
Header->GetChunkStatePtr(WriteChunkIdx).StartReading(ReaderIdx);
}
return ReaderIdx;
}
}
}
}
}
return -1;
}
void AddReaderRef(int ReaderIdx)
{
FReaderStatePtr ReaderStatePtr = Header->GetReaderStatePtr(ReaderIdx);
for (; ; )
{
FReaderState ReaderState = ReaderStatePtr.Get();
check(ReaderState.GetRefCount() > 0);
if (ReaderStatePtr.TryUpdate(ReaderState, FReaderState(ReaderState.GetChunkIdx(), ReaderState.GetOffset(), ReaderState.GetRefCount() + 1, ReaderState.IsDetached())))
{
break;
}
}
}
void ReleaseReaderRef(int ReaderIdx)
{
FReaderStatePtr ReaderStatePtr = Header->GetReaderStatePtr(ReaderIdx);
for (; ; )
{
FReaderState ReaderState = ReaderStatePtr.Get();
check(ReaderState.GetRefCount() > 0);
if (ReaderState.GetRefCount() == 1)
{
for (int Idx = 0; Idx < Header->NumChunks; Idx++)
{
Header->GetChunkStatePtr(Idx).FinishReading(ReaderIdx);
}
}
if (ReaderStatePtr.TryUpdate(ReaderState, FReaderState(ReaderState.GetChunkIdx(), ReaderState.GetOffset(), ReaderState.GetRefCount() - 1, ReaderState.IsDetached())))
{
break;
}
}
}
void CreateWriter()
{
FWriterStatePtr WriterStatePtr = Header->GetWriterStatePtr();
for (; ; )
{
FWriterState WriterState = WriterStatePtr.Get();
int WriterStateRefCount = WriterState.GetRefCount();
checkf(WriterStateRefCount == 0, TEXT("Attempting to CreateWriter with a non-zero WriterState refcount: %d"), WriterStateRefCount);
if (WriterStatePtr.TryUpdate(WriterState, FWriterState(WriterState.GetChunkIdx(), WriterState.GetReaderFlags(), 1, WriterState.HasWrapped())))
{
FChunkStatePtr ChunkStatePtr = Header->GetChunkStatePtr(WriterState.GetChunkIdx());
for (; ; )
{
FChunkState ChunkState = ChunkStatePtr.Get();
if (ChunkStatePtr.TryUpdate(ChunkState, FChunkState(EWriteState::Writing, ChunkState.GetReaderFlags(), ChunkState.GetLength())))
{
break;
}
}
break;
}
}
}
void AddWriterRef()
{
FWriterStatePtr WriterStatePtr = Header->GetWriterStatePtr();
for (; ; )
{
FWriterState WriterState = WriterStatePtr.Get();
check(WriterState.GetRefCount() > 0);
if (WriterStatePtr.TryUpdate(WriterState, FWriterState(WriterState.GetChunkIdx(), WriterState.GetReaderFlags(), WriterState.GetRefCount() + 1, WriterState.HasWrapped())))
{
break;
}
}
}
void ReleaseWriterRef()
{
FWriterStatePtr writerStatePtr = Header->GetWriterStatePtr();
for (; ; )
{
FWriterState writerState = writerStatePtr.Get();
check(writerState.GetRefCount() > 0);
if (writerState.GetRefCount() == 1)
{
MarkComplete();
}
if (writerStatePtr.TryUpdate(writerState, FWriterState(writerState.GetChunkIdx(), writerState.GetReaderFlags(), writerState.GetRefCount() - 1, writerState.HasWrapped())))
{
break;
}
}
}
bool IsComplete(int ReaderIdx) const
{
FReaderState ReaderState = Header->GetReaderStatePtr(ReaderIdx).Get();
if (ReaderState.IsDetached())
{
return true;
}
FChunkState ChunkState = Header->GetChunkStatePtr(ReaderState.GetChunkIdx()).Get();
return ChunkState.GetWriteState() == EWriteState::Complete && ReaderState.GetOffset() == ChunkState.GetLength();
}
void DetachReader(int ReaderIdx)
{
FReaderStatePtr ReaderStatePtr = Header->GetReaderStatePtr(ReaderIdx);
for (; ; )
{
FReaderState ReaderState = ReaderStatePtr.Get();
if (ReaderStatePtr.TryUpdate(ReaderState, FReaderState(ReaderState.GetChunkIdx(), ReaderState.GetOffset(), ReaderState.GetRefCount(), true)))
{
ReaderEvents[ReaderIdx].Signal();
break;
}
}
}
void AdvanceReadPosition(int ReaderIdx, size_t Offset)
{
FReaderStatePtr ReaderStatePtr = Header->GetReaderStatePtr(ReaderIdx);
for (; ; )
{
FReaderState ReaderState = ReaderStatePtr.Get();
if (ReaderStatePtr.TryUpdate(ReaderState, FReaderState(ReaderState.GetChunkIdx(), ReaderState.GetOffset() + (int)Offset, ReaderState.GetRefCount(), ReaderState.IsDetached())))
{
ReaderEvents[ReaderIdx].Signal();
break;
}
}
}
size_t GetMaxReadSize(int ReaderIdx) const
{
FReaderState ReaderState = Header->GetReaderStatePtr(ReaderIdx).Get();
if (ReaderState.IsDetached())
{
return 0;
}
FChunkStatePtr ChunkStatePtr = Header->GetChunkStatePtr(ReaderState.GetChunkIdx());
FChunkState ChunkState = ChunkStatePtr.Get();
if (ChunkState.HasReaderFlag(ReaderIdx))
{
return ChunkState.GetLength() - ReaderState.GetOffset();
}
else
{
return 0;
}
}
const unsigned char* WaitToRead(int ReaderIdx, size_t MinSize, int TimeoutMs, bool* bOutHasTimedOut)
{
for (; ; )
{
FReaderStatePtr ReaderStatePtr = Header->GetReaderStatePtr(ReaderIdx);
FReaderState ReaderState = ReaderStatePtr.Get();
if (ReaderState.IsDetached())
{
return nullptr;
}
FChunkStatePtr ChunkStatePtr = Header->GetChunkStatePtr(ReaderState.GetChunkIdx());
FChunkState ChunkState = ChunkStatePtr.Get();
if (!ChunkState.HasReaderFlag(ReaderIdx))
{
// Wait until the current chunk is readable
if (!ReaderEvents[ReaderIdx].Wait(TimeoutMs))
{
if (bOutHasTimedOut)
{
*bOutHasTimedOut = true;
}
return nullptr;
}
}
else if (ReaderState.GetOffset() + MinSize <= ChunkState.GetLength())
{
// We have enough data in the chunk to be able to read a message
return ChunkPtrs[ReaderState.GetChunkIdx()] + ReaderState.GetOffset();
}
else if (ChunkState.GetWriteState() == EWriteState::Writing)
{
// Wait until there is more data in the chunk
if (!ReaderEvents[ReaderIdx].Wait(TimeoutMs))
{
if (bOutHasTimedOut)
{
*bOutHasTimedOut = true;
}
return nullptr;
}
}
else if (ReaderState.GetOffset() < ChunkState.GetLength() || ChunkState.GetWriteState() == EWriteState::Complete)
{
// Cannot read the requested amount of data from this chunk.
return nullptr;
}
else if (ChunkState.GetWriteState() == EWriteState::MovedToNext)
{
// Move to the next chunk
ChunkStatePtr.FinishReading(ReaderIdx);
WriterEvent.Signal();
int NextChunkIdx = ReaderState.GetChunkIdx() + 1;
if (NextChunkIdx == Header->NumChunks)
{
NextChunkIdx = 0;
}
ReaderStatePtr.TryUpdate(ReaderState, FReaderState(NextChunkIdx, 0, ReaderState.GetRefCount(), ReaderState.IsDetached()));
}
else
{
check(false);
}
}
}
bool MarkComplete()
{
FWriterState WriterState = Header->GetWriterStatePtr().Get();
FChunkStatePtr ChunkStatePtr = Header->GetChunkStatePtr(WriterState.GetChunkIdx());
if (ChunkStatePtr.Get().GetWriteState() != EWriteState::Complete)
{
ChunkStatePtr.MarkComplete();
SetAllReadEvents();
return true;
}
return false;
}
void AdvanceWritePosition(size_t Size)
{
if (Size > 0)
{
FWriterState WriterState = Header->GetWriterStatePtr().Get();
FChunkStatePtr ChunkStatePtr = Header->GetChunkStatePtr(WriterState.GetChunkIdx());
FChunkState ChunkState = ChunkStatePtr.Get();
check(ChunkState.GetWriteState() == EWriteState::Writing);
ChunkStatePtr.Append(Size);
SetAllReadEvents();
}
}
size_t GetMaxWriteSize() const
{
FWriterState WriterState = Header->GetWriterStatePtr().Get();
FChunkState ChunkState = Header->GetChunkStatePtr(WriterState.GetChunkIdx()).Get();
if (ChunkState.GetWriteState() == EWriteState::Writing)
{
return Header->ChunkLength - ChunkState.GetLength();
}
else
{
return 0;
}
}
size_t Write(const void* Buffer, size_t MaxSize, int TimeoutMs)
{
unsigned char* SendData = WaitToWrite(1, TimeoutMs);
if (SendData == nullptr)
{
return 0;
}
size_t SendSize = GetMaxWriteSize();
if (MaxSize < SendSize)
{
SendSize = MaxSize;
}
memcpy(SendData, Buffer, SendSize);
AdvanceWritePosition(SendSize);
return SendSize;
}
unsigned char* WaitToWrite(size_t MinSize, int TimeoutMs)
{
check(MinSize <= Header->ChunkLength);
// Get the current chunk we're writing to
FWriterState WriterState = Header->GetWriterStatePtr().Get();
int WriteChunkIdx = WriterState.GetChunkIdx();
FChunkStatePtr WriteChunkStatePtr = Header->GetChunkStatePtr(WriteChunkIdx);
// Check if we can append to this chunk
FChunkState ChunkState = WriteChunkStatePtr.Get();
if (ChunkState.GetWriteState() == EWriteState::Writing)
{
int Length = ChunkState.GetLength();
if (Length + MinSize <= Header->ChunkLength)
{
return ChunkPtrs[WriteChunkIdx] + Length;
}
WriteChunkStatePtr.FinishWriting(); // STATE CHANGE
SetAllReadEvents();
}
if (ChunkState.GetWriteState() == EWriteState::Complete)
{
return nullptr;
}
// Otherwise get the next chunk to write to
int NextWriteChunkIdx = WriteChunkIdx + 1;
if (NextWriteChunkIdx == Header->NumChunks)
{
NextWriteChunkIdx = 0;
}
// Wait until all readers have finished with the chunk, and we can update the writer to match
FChunkStatePtr NextWriteChunkStatePtr = Header->GetChunkStatePtr(NextWriteChunkIdx);
for (;;)
{
FChunkState NextWriteChunkState = NextWriteChunkStatePtr.Get();
if (NextWriteChunkState.GetReaderFlags() != 0)
{
if (!WriterEvent.Wait(TimeoutMs))
{
return nullptr;
}
}
else if (NextWriteChunkStatePtr.TryUpdate(NextWriteChunkState, FChunkState(EWriteState::Writing, WriterState.GetReaderFlags(), 0)))
{
FWriterStatePtr WriterStatePtr = Header->GetWriterStatePtr();
if (WriterStatePtr.TryUpdate(WriterState, FWriterState(NextWriteChunkIdx, WriterState.GetReaderFlags(), WriterState.GetRefCount(), NextWriteChunkIdx == 0)))
{
break;
}
else
{
WriterState = WriterStatePtr.Get();
}
}
}
return ChunkPtrs[NextWriteChunkIdx];
}
private:
FComputeMemoryMappedFile MemoryMappedFile;
static void InitShared(FComputeBufferDetail* Detail)
{
FHeader* Header = Detail->Header;
unsigned char* NextPtr = (unsigned char*)(Header + 1);
for (int ChunkIdx = 0; ChunkIdx < Header->NumChunks; ChunkIdx++)
{
Detail->ChunkPtrs[ChunkIdx] = NextPtr;
NextPtr += Header->ChunkLength;
}
}
void SetAllReadEvents()
{
for (int Idx = 0; Idx < Header->NumReaders; Idx++)
{
ReaderEvents[Idx].Signal();
}
}
};
//// FComputeBuffer /////
FComputeBuffer::FComputeBuffer()
: Detail(nullptr)
{
}
FComputeBuffer::FComputeBuffer(const FComputeBuffer& Other)
: Detail(Other.Detail)
{
if (Detail)
{
Detail->AddRef();
}
}
FComputeBuffer::FComputeBuffer(FComputeBuffer&& Other) noexcept
: Detail(Other.Detail)
{
Other.Detail = nullptr;
}
FComputeBuffer::~FComputeBuffer()
{
Close();
}
FComputeBuffer& FComputeBuffer::operator=(const FComputeBuffer& Other)
{
Close();
if (Other.Detail)
{
Detail = Other.Detail;
Detail->AddRef();
}
return *this;
}
FComputeBuffer& FComputeBuffer::operator=(FComputeBuffer&& Other) noexcept
{
Close();
Detail = Other.Detail;
Other.Detail = nullptr;
return *this;
}
bool FComputeBuffer::CreateNew(const FParams& Params)
{
Close();
TUniquePtr<FComputeBufferDetail> Buffer = FComputeBufferDetail::CreateNew(Params);
if (Buffer)
{
Detail = Buffer.Release();
return true;
}
return false;
}
bool FComputeBuffer::OpenExisting(const char* Name)
{
Close();
TUniquePtr<FComputeBufferDetail> Buffer = FComputeBufferDetail::OpenExisting(Name);
if (Buffer)
{
Detail = Buffer.Release();
return true;
}
return false;
}
void FComputeBuffer::Close()
{
if (Detail != nullptr)
{
Detail->Release();
Detail = nullptr;
}
}
FComputeBufferReader FComputeBuffer::CreateReader()
{
int ReaderIdx = Detail->CreateReader();
Detail->AddRef();
return FComputeBufferReader(Detail, ReaderIdx);
}
FComputeBufferWriter FComputeBuffer::CreateWriter()
{
Detail->CreateWriter();
Detail->AddRef();
return FComputeBufferWriter(Detail);
}
const char* FComputeBuffer::GetName() const
{
return Detail->Name;
}
//// FComputeBufferReader /////
FComputeBufferReader::FComputeBufferReader()
: Detail(nullptr)
, ReaderIdx(0)
{
}
FComputeBufferReader::FComputeBufferReader(const FComputeBufferReader& Other)
: Detail(Other.Detail)
, ReaderIdx(Other.ReaderIdx)
{
if (Detail)
{
Detail->AddRef();
Detail->AddReaderRef(ReaderIdx);
}
}
FComputeBufferReader::FComputeBufferReader(FComputeBufferReader&& Other) noexcept
: Detail(Other.Detail)
, ReaderIdx(Other.ReaderIdx)
{
Other.Detail = nullptr;
}
FComputeBufferReader::~FComputeBufferReader()
{
Close();
}
FComputeBufferReader& FComputeBufferReader::operator=(const FComputeBufferReader& Other)
{
Close();
if (Other.Detail)
{
Detail = Other.Detail;
ReaderIdx = Other.ReaderIdx;
Detail->AddRef();
Detail->AddReaderRef(ReaderIdx);
}
return *this;
}
FComputeBufferReader& FComputeBufferReader::operator=(FComputeBufferReader&& Other) noexcept
{
Close();
Detail = Other.Detail;
ReaderIdx = Other.ReaderIdx;
Other.Detail = nullptr;
return *this;
}
void FComputeBufferReader::Close()
{
if (Detail != nullptr)
{
Detail->ReleaseReaderRef(ReaderIdx);
Detail->Release();
Detail = nullptr;
}
}
void FComputeBufferReader::Detach()
{
check(Detail);
Detail->DetachReader(ReaderIdx);
}
bool FComputeBufferReader::IsComplete() const
{
check(Detail);
return Detail->IsComplete(ReaderIdx);
}
void FComputeBufferReader::AdvanceReadPosition(size_t Size)
{
check(Detail);
Detail->AdvanceReadPosition(ReaderIdx, Size);
}
size_t FComputeBufferReader::GetMaxReadSize() const
{
check(Detail);
return Detail->GetMaxReadSize(ReaderIdx);
}
size_t FComputeBufferReader::Read(void* Buffer, size_t MaxSize, int TimeoutMs, bool* bOutTimedOut)
{
const unsigned char* RecvData = WaitToRead(1, TimeoutMs, bOutTimedOut);
if (RecvData == nullptr)
{
return 0;
}
size_t RecvSize = GetMaxReadSize();
if (MaxSize < RecvSize)
{
RecvSize = MaxSize;
}
memcpy(Buffer, RecvData, RecvSize);
AdvanceReadPosition(RecvSize);
return RecvSize;
}
const unsigned char* FComputeBufferReader::WaitToRead(size_t MinSize, int TimeoutMs, bool* bOutTimedOut)
{
check(Detail);
return Detail->WaitToRead(ReaderIdx, MinSize, TimeoutMs, bOutTimedOut);
}
FComputeBufferReader::FComputeBufferReader(FComputeBufferDetail* InDetail, int InReaderIdx)
: Detail(InDetail)
, ReaderIdx(InReaderIdx)
{
}
const char* FComputeBufferReader::GetName() const
{
check(Detail);
return Detail->Name;
}
//// FComputeBufferWriter /////
FComputeBufferWriter::FComputeBufferWriter()
: Detail(nullptr)
{
}
FComputeBufferWriter::FComputeBufferWriter(const FComputeBufferWriter& Other)
: Detail(Other.Detail)
{
if (Detail)
{
Detail->AddRef();
Detail->AddWriterRef();
}
}
FComputeBufferWriter::FComputeBufferWriter(FComputeBufferWriter&& Other) noexcept
: Detail(Other.Detail)
{
Other.Detail = nullptr;
}
FComputeBufferWriter::~FComputeBufferWriter()
{
Close();
}
FComputeBufferWriter& FComputeBufferWriter::operator=(const FComputeBufferWriter& Other)
{
Close();
if (Other.Detail)
{
Detail = Other.Detail;
Detail->AddWriterRef();
Detail->AddRef();
}
return *this;
}
FComputeBufferWriter& FComputeBufferWriter::operator=(FComputeBufferWriter&& Other) noexcept
{
Close();
Detail = Other.Detail;
Other.Detail = nullptr;
return *this;
}
void FComputeBufferWriter::Close()
{
if (Detail != nullptr)
{
Detail->ReleaseWriterRef();
Detail->Release();
Detail = nullptr;
}
}
void FComputeBufferWriter::MarkComplete()
{
Detail->MarkComplete();
}
void FComputeBufferWriter::AdvanceWritePosition(size_t Size)
{
Detail->AdvanceWritePosition(Size);
}
size_t FComputeBufferWriter::GetMaxWriteSize() const
{
return Detail->GetMaxWriteSize();
}
size_t FComputeBufferWriter::GetChunkMaxLength() const
{
return Detail->Header->ChunkLength;
}
size_t FComputeBufferWriter::Write(const void* Buffer, size_t MaxSize, int TimeoutMs)
{
unsigned char* SendData = WaitToWrite(1, TimeoutMs);
if (SendData == nullptr)
{
return 0;
}
size_t SendSize = GetMaxWriteSize();
if (MaxSize < SendSize)
{
SendSize = MaxSize;
}
memcpy(SendData, Buffer, SendSize);
AdvanceWritePosition(SendSize);
return SendSize;
}
unsigned char* FComputeBufferWriter::WaitToWrite(size_t MinSize, int TimeoutMs)
{
return Detail->WaitToWrite(MinSize, TimeoutMs);
}
FComputeBufferWriter::FComputeBufferWriter(FComputeBufferDetail* InDetail)
: Detail(InDetail)
{
}
const char* FComputeBufferWriter::GetName() const
{
return Detail->Name;
}