Skip to content

perf(screenshot): Avoid redundant buffers and copies in screenshot uploads #635

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tintinhamans wants to merge 1 commit into
base: main
Choose a base branch
from
+88 −121
Open

perf(screenshot): Avoid redundant buffers and copies in screenshot uploads #635

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions GeneralsMD/Code/GameEngine/Include/GameNetwork/GeneralsOnline/OnlineServices_Init.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -586,10 +586,10 @@ class NGMP_OnlineServicesManager
ServiceConfig& GetServiceConfig() { return m_ServiceConfig; }

public:
void CacheScreenshotBytes_StartMatch(std::vector<uint8_t>& vecData)
void CacheScreenshotBytes_StartMatch(std::vector<uint8_t> vecData)
{
std::scoped_lock<std::mutex> ssLock(m_ScreenshotMutex);
m_vecCachedScreenshotBytes_MatchStart = vecData;
m_vecCachedScreenshotBytes_MatchStart = std::move(vecData);
}

void CacheScreenshotBytes_EndMatch(uint64_t matchID, std::vector<uint8_t> data);
Expand Down
2 changes: 1 addition & 1 deletion GeneralsMD/Code/GameEngine/Source/GameNetwork/GeneralsOnline/HTTP/HTTPManager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ void HTTPManager::SendS3PUTRequest(const char* szURI, EIPProtocolVersion protove

HTTPRequest* pRequest = PlatformCreateRequest(EHTTPVerb::HTTP_VERB_PUT, protover, szURI, inHeaders, completionCallback, progressCallback, timeoutMS);
pRequest->DisableServiceAuth();
pRequest->SetPostDataBuffer(vecBuffer);
pRequest->SetPostDataBuffer(std::move(vecBuffer));

m_vecRequestsPendingStart.push_back(pRequest);
}
Expand Down
203 changes: 85 additions & 118 deletions GeneralsMD/Code/GameEngine/Source/GameNetwork/GeneralsOnline/OnlineServices_Init.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ void NGMP_OnlineServicesManager::CaptureScreenshotForProbe(EScreenshotType scree
// certain screenshots require caching for later upload when we have a valid match id, so we store them
if (screenshotType == EScreenshotType::SCREENSHOT_TYPE_LOADSCREEN)
{
NGMP_OnlineServicesManager::GetInstance()->CacheScreenshotBytes_StartMatch(vecData);
NGMP_OnlineServicesManager::GetInstance()->CacheScreenshotBytes_StartMatch(std::move(vecData));
}
else if (screenshotType == EScreenshotType::SCREENSHOT_TYPE_SCORESCREEN)
{
Expand All @@ -143,7 +143,7 @@ void NGMP_OnlineServicesManager::CaptureScreenshotForProbe(EScreenshotType scree
newEntry.vecBytes = std::move(vecData);
newEntry.strSignedURI = strURI;
newEntry.screenshotType = screenshotType;
m_vecGuardedSSData.push_back(newEntry);
m_vecGuardedSSData.push_back(std::move(newEntry));
}
});
}
Expand Down Expand Up @@ -579,7 +579,6 @@ void NGMP_OnlineServicesManager::CaptureScreenshot(bool bResizeForTransmit, std:
}

SurfaceClass* surface = DX8Wrapper::_Get_DX8_Back_Buffer();
LPDIRECT3DSURFACE8 surf = nullptr;
SurfaceClass* surfaceCopy = nullptr;
void* pBits = nullptr;
IDirect3DSurface8* pDXsurf = nullptr;
Expand All @@ -599,138 +598,106 @@ void NGMP_OnlineServicesManager::CaptureScreenshot(bool bResizeForTransmit, std:
{
DX8Wrapper::_Copy_DX8_Rects(surface->Peek_D3D_Surface(), NULL, 0, surfaceCopy->Peek_D3D_Surface(), NULL);

HRESULT hr;
// gather all our data
int pitch = 0;
pBits = surfaceCopy->Lock(&pitch);

D3DDISPLAYMODE mode;
if (SUCCEEDED(hr = DX8Wrapper::_Get_D3D_Device8()->GetDisplayMode(&mode)))
if (pBits != nullptr)
{
if (SUCCEEDED(hr = DX8Wrapper::_Get_D3D_Device8()->CreateImageSurface(mode.Width, mode.Height,
D3DFMT_A8R8G8B8, &surf)))
{
if (SUCCEEDED(hr = DX8Wrapper::_Get_D3D_Device8()->GetFrontBuffer(surf)))
int width = surfaceDesc.Width;
int height = surfaceDesc.Height;

// Copy pixel data into an owned buffer before spawning the thread so
// the surface can be safely unlocked on the main thread immediately after.
const uint8_t* pSurfaceBytes = static_cast<const uint8_t*>(pBits);
std::vector<uint8_t> pixelData(pSurfaceBytes, pSurfaceBytes + height * pitch);

// process on thread - track the thread so we can join it during shutdown
std::thread* pNewThread = new std::thread([cbOnDataAvailable, width, height, pixelData = std::move(pixelData), pitch, bResizeForTransmit]() mutable
{
// gather all our data
int pitch = 0;
pBits = surfaceCopy->Lock(&pitch);
CHECK_WORKER_THREAD;

if (pBits != nullptr)
{
int width = surfaceDesc.Width;
int height = surfaceDesc.Height;

// Copy pixel data into an owned buffer before spawning the thread so
// the surface can be safely unlocked on the main thread immediately after.
std::vector<uint8_t> pixelData(height * pitch);
memcpy(pixelData.data(), pBits, height * pitch);

// process on thread - track the thread so we can join it during shutdown
std::thread* pNewThread = new std::thread([cbOnDataAvailable, width, height, pixelData = std::move(pixelData), pitch, bResizeForTransmit]()
{
CHECK_WORKER_THREAD;

unsigned char* rgbData = new unsigned char[width * height * 3];

std::vector<unsigned char> vecData;

int finalWidth = width;
int finalHeight = height;

for (int y = 0; y < height; ++y) {
const uint8_t* row = pixelData.data() + y * pitch;
int rowOffset = y * width * 3;
int srcOffset = 0;
for (int x = 0; x < width; ++x, srcOffset += 4)
{
int dstIndex = rowOffset + x * 3;
rgbData[dstIndex + 0] = row[srcOffset + 2]; // R
rgbData[dstIndex + 1] = row[srcOffset + 1]; // G
rgbData[dstIndex + 2] = row[srcOffset + 0]; // B
}
}

// resize
unsigned char* pBufferToWrite = rgbData;
if (bResizeForTransmit)
{
ServiceConfig& serviceConf = NGMP_OnlineServicesManager::GetInstance()->GetServiceConfig();
int new_width = serviceConf.screenshot_width;
int new_height = serviceConf.screenshot_height;
int channels = 3;
unsigned char* resized = new unsigned char[new_width * new_height * channels];

stbir_resize_uint8(rgbData, width, height, 0,
resized, new_width, new_height, 0,
channels
);

// update data
finalWidth = new_width;
finalHeight = new_height;
pBufferToWrite = resized;
}
// end resize

stbi_write_jpg_to_func([](void* context, void* data, int size)
{
std::vector<unsigned char>* buffer = static_cast<std::vector<unsigned char>*>(context);
buffer->insert(buffer->end(), (unsigned char*)data, (unsigned char*)data + size);
}, &vecData, finalWidth, finalHeight, 3, pBufferToWrite, bResizeForTransmit ? 0 : 90);

// cleanup
if (bResizeForTransmit)
{
delete[] pBufferToWrite; // This is 'resized'
pBufferToWrite = nullptr;
}

delete[] rgbData;
rgbData = nullptr;

// invoke cb
if (cbOnDataAvailable != nullptr)
{
cbOnDataAvailable(vecData);
}
}
);
const int channels = 3;
int finalWidth = width;
int finalHeight = height;

// Store the thread so we can join it during shutdown
// SECURITY FIX: Capture manager pointer before spawning thread to avoid TOCTOU race
NGMP_OnlineServicesManager* pMgr = NGMP_OnlineServicesManager::GetInstance();
if (pMgr != nullptr)
// Pack the BGRA rows to tight RGB in place. The write position never passes the read position.
uint8_t* pixels = pixelData.data();
for (int y = 0; y < height; ++y)
{
std::scoped_lock<std::mutex> lock(pMgr->m_mutexScreenshotThreads);
pMgr->m_vecScreenshotThreads.push_back(pNewThread);
const uint8_t* src = pixels + y * pitch;
uint8_t* dst = pixels + y * width * channels;
for (int x = 0; x < width; ++x, src += 4, dst += channels)
{
const uint8_t blue = src[0];
dst[0] = src[2]; // R
dst[1] = src[1]; // G
dst[2] = blue; // B
}
}
else

// resize
std::vector<uint8_t> resized;
if (bResizeForTransmit)
{
// Manager was destroyed, cannot store thread. Thread will leak but won't crash.
NetworkLog(ELogVerbosity::LOG_RELEASE, "[Screenshot] Manager destroyed before thread could be registered");
ServiceConfig& serviceConf = NGMP_OnlineServicesManager::GetInstance()->GetServiceConfig();
finalWidth = serviceConf.screenshot_width;
finalHeight = serviceConf.screenshot_height;
resized.resize(finalWidth * finalHeight * channels);

stbir_resize_uint8(pixels, width, height, 0,
resized.data(), finalWidth, finalHeight, 0,
channels
);

// the capture is no longer needed
std::vector<uint8_t>().swap(pixelData);
}
// end resize

// room for a typical JPEG, so the write callback rarely has to grow it
std::vector<unsigned char> vecData;
vecData.reserve(256 * 1024);

stbi_write_jpg_to_func([](void* context, void* data, int size)
{
std::vector<unsigned char>* buffer = static_cast<std::vector<unsigned char>*>(context);
buffer->insert(buffer->end(), (unsigned char*)data, (unsigned char*)data + size);
}, &vecData, finalWidth, finalHeight, channels, bResizeForTransmit ? resized.data() : pixelData.data(), bResizeForTransmit ? 0 : 90);

bSucceeded = true;
// invoke cb
if (cbOnDataAvailable != nullptr)
{
cbOnDataAvailable(std::move(vecData));
}
}
);

// Store the thread so we can join it during shutdown
// SECURITY FIX: Capture manager pointer before spawning thread to avoid TOCTOU race
NGMP_OnlineServicesManager* pMgr = NGMP_OnlineServicesManager::GetInstance();
if (pMgr != nullptr)
{
std::scoped_lock<std::mutex> lock(pMgr->m_mutexScreenshotThreads);
pMgr->m_vecScreenshotThreads.push_back(pNewThread);
}
else
{
// Manager was destroyed, cannot store thread. Thread will leak but won't crash.
NetworkLog(ELogVerbosity::LOG_RELEASE, "[Screenshot] Manager destroyed before thread could be registered");
}

bSucceeded = true;
}
}
}
}

// clean everything up, whether we succeeded or not

// release the image surface
if (surf != nullptr)
{
surf->Release();
//delete surf;
surf = nullptr;
}

// unlock
if (surface != nullptr)
{
surface->Unlock();
surface->Release_Ref();
surface = nullptr;
}
Expand Down Expand Up @@ -954,7 +921,7 @@ void NGMP_OnlineServicesManager::Tick()


// screenshot types that already have a presigned URL
for (S3ScreenshotEntry screenshotEntry : m_vecGuardedSSData)
for (S3ScreenshotEntry& screenshotEntry : m_vecGuardedSSData)
{
// NOTE: Screenshot types start and end of match are captures and cached in memory until the server tells us where to upload them, so we need to wait for the upload URI

Expand All @@ -966,7 +933,7 @@ void NGMP_OnlineServicesManager::Tick()
{
std::map<std::string, std::string> mapHeaders;
mapHeaders["Content-Type"] = "image/jpeg";
NGMP_OnlineServicesManager::GetInstance()->GetHTTPManager()->SendS3PUTRequest(screenshotEntry.strSignedURI.c_str(), EIPProtocolVersion::DONT_CARE, mapHeaders, screenshotEntry.vecBytes, [=](bool bSuccess, int statusCode, std::string strBody, HTTPRequest* pReq)
NGMP_OnlineServicesManager::GetInstance()->GetHTTPManager()->SendS3PUTRequest(screenshotEntry.strSignedURI.c_str(), EIPProtocolVersion::DONT_CARE, mapHeaders, std::move(screenshotEntry.vecBytes), [=](bool bSuccess, int statusCode, std::string strBody, HTTPRequest* pReq)
{
#if _DEBUG
if (statusCode != 200)
Expand All @@ -990,9 +957,9 @@ void NGMP_OnlineServicesManager::Tick()
// queue it
S3ScreenshotEntry newEntry;
newEntry.screenshotType = EScreenshotType::SCREENSHOT_TYPE_LOADSCREEN;
newEntry.vecBytes = m_vecCachedScreenshotBytes_MatchStart;
newEntry.strSignedURI = m_strCachedScreenshot_MatchStart_S3URI;
m_vecGuardedSSData.push_back(newEntry);
newEntry.vecBytes = std::move(m_vecCachedScreenshotBytes_MatchStart);
newEntry.strSignedURI = std::move(m_strCachedScreenshot_MatchStart_S3URI);
m_vecGuardedSSData.push_back(std::move(newEntry));

// clear data
m_vecCachedScreenshotBytes_MatchStart = std::vector<uint8_t>();
Expand Down Expand Up @@ -1023,7 +990,7 @@ void NGMP_OnlineServicesManager::Tick()
// do the upload
std::map<std::string, std::string> mapHeaders;
mapHeaders["Content-Type"] = "application/octet-stream";
NGMP_OnlineServicesManager::GetInstance()->GetHTTPManager()->SendS3PUTRequest(m_cachedReplayUpload.signedURI.c_str(), EIPProtocolVersion::DONT_CARE, mapHeaders, m_cachedReplayUpload.bytes, [=](bool bSuccess, int statusCode, std::string strBody, HTTPRequest* pReq)
NGMP_OnlineServicesManager::GetInstance()->GetHTTPManager()->SendS3PUTRequest(m_cachedReplayUpload.signedURI.c_str(), EIPProtocolVersion::DONT_CARE, mapHeaders, std::move(m_cachedReplayUpload.bytes), [=](bool bSuccess, int statusCode, std::string strBody, HTTPRequest* pReq)
{
#if _DEBUG
if (statusCode != 200)
Expand Down
Loading