DX11 GPU memory leak on OpenSharedResource

I'm trying to figure why i'm getting GPU memory leak on this code.

I checked using System Informer, at "Gpu Dedicated bytes" the memory usage increases on each loop iteration.

I thought it was a missing call to CloseHandle(hDxSurface);, but this throws an exception.

I'm not sure what else is causing the leak.

#include "stdafx.h"
#include <Windows.h>
#include <chrono>
#include <thread>
#include <atomic>
#include <memory>
#include <d3d11.h>
#include <dxgi1_2.h>
#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "dxgi.lib")



typedef HRESULT(WINAPI* PFN_DwmDxGetWindowSharedSurface)(HWND, HANDLE*, LUID*, ULONG*, ULONG*, ULONGLONG*);
PFN_DwmDxGetWindowSharedSurface DwmGetDxSharedSurface = nullptr;

bool getDwmGetDxSharedSurface()
{
    HMODULE hUser32 = LoadLibraryA("user32.dll");
    if (hUser32 != NULL)
        DwmGetDxSharedSurface = (PFN_DwmDxGetWindowSharedSurface)GetProcAddress(hUser32, "DwmGetDxSharedSurface");
    return DwmGetDxSharedSurface != nullptr;
}



class WindowCapture
{
public:    
    // DirectX resources (cached and reused)
    IDXGIFactory1*       m_pFactory = nullptr;
    IDXGIAdapter*        m_pAdapter = nullptr;
    ID3D11DeviceContext* m_pContext = nullptr;
    ID3D11Device*        m_pDevice  = nullptr;
    // Target window
    HWND m_hwnd = nullptr;
    
    // **PERFORMANCE OPTIMIZATION**: Double-buffered staging textures
    ID3D11Texture2D* m_pStagingTexture[2] = {nullptr, nullptr};
    int m_currentStagingIndex = 0;
    UINT m_lastWidth          = 0;
    UINT m_lastHeight         = 0;
    DXGI_FORMAT m_lastFormat  = DXGI_FORMAT_UNKNOWN;
    
    // DIB Section for direct pixel access (cached)
    HBITMAP m_hDIBSection = nullptr;
    void* m_pDIBBits      = nullptr;
    UINT m_dibWidth       = 0;
    UINT m_dibHeight      = 0;

    WindowCapture() {};
    bool initialize(HWND hwnd);
    bool captureToPixelBuffer(void** ppPixels, UINT* pWidth, UINT* pHeight, UINT* pPitch);
    bool captureToFile(const wchar_t* filename);
    bool initializeDirectX();
    bool createOrUpdateStagingTextures(UINT width, UINT height, DXGI_FORMAT format);
    bool createOrUpdateDIBSection(UINT width, UINT height);
    bool copyTextureDataToDIB(const D3D11_MAPPED_SUBRESOURCE& mapped, UINT width, UINT height);
};



bool WindowCapture::initialize(HWND hwnd)
{
    if (!hwnd || !IsWindow(hwnd))
        return false;
    
    m_hwnd = hwnd;
    if (!getDwmGetDxSharedSurface())
    {
        $err, "Failed to get DwmGetDxSharedSurface function";
        return false;
    }

    HRESULT hr = CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED);
    if (FAILED(hr) && hr != RPC_E_CHANGED_MODE)
    {
    	$err, "Failed to initialize COM", hr;
        return false;
    }

    if (!initializeDirectX())
        return false;
    return true;
}



bool WindowCapture::initializeDirectX()
{
    HRESULT hr = CreateDXGIFactory1(__uuidof(IDXGIFactory1), (void**)(&m_pFactory));
    if (FAILED(hr) || !m_pFactory)
    {
        $err, "CreateDXGIFactory1 failed", hr;
        return false;
    }

    m_pFactory->EnumAdapters(0, &m_pAdapter);
    if (!m_pAdapter)
    {
        $err, "Failed to enumerate adapters";
        m_pFactory->Release();
        return false;
	}

    const D3D_FEATURE_LEVEL featureLevels[] = 
    {
        D3D_FEATURE_LEVEL_11_0,
        D3D_FEATURE_LEVEL_10_1,
        D3D_FEATURE_LEVEL_10_0,
        D3D_FEATURE_LEVEL_9_3,
        D3D_FEATURE_LEVEL_9_2,
        D3D_FEATURE_LEVEL_9_1
    };

    hr = D3D11CreateDevice(m_pAdapter, D3D_DRIVER_TYPE_UNKNOWN, NULL, D3D11_CREATE_DEVICE_BGRA_SUPPORT, featureLevels, 6, D3D11_SDK_VERSION, &m_pDevice, NULL, &m_pContext);
    if (FAILED(hr) || !m_pDevice)
    {
        $err, "D3D11CreateDevice failed", hr;
        m_pFactory->Release();
        return false;
    }

    return true;
}



bool WindowCapture::createOrUpdateStagingTextures(UINT width, UINT height, DXGI_FORMAT format)
{
    // Reuse existing staging textures if dimensions and format match
    if (m_pStagingTexture[0] && m_lastWidth == width && m_lastHeight == height && m_lastFormat == format)
        return true;
    
    // Release old staging textures
    for (int i = 0; i < 2; ++i)
    {
        if (m_pStagingTexture[i])
        {
            m_pStagingTexture[i]->Release();
            m_pStagingTexture[i] = nullptr;
        }
    }
    
    // Create new staging textures (double-buffered)
    D3D11_TEXTURE2D_DESC stagingDesc = {};
    stagingDesc.Width              = width;
    stagingDesc.Height             = height;
    stagingDesc.MipLevels          = 1;
    stagingDesc.ArraySize          = 1;
    stagingDesc.Format             = format;
    stagingDesc.SampleDesc.Count   = 1;
    stagingDesc.SampleDesc.Quality = 0;
    stagingDesc.Usage              = D3D11_USAGE_STAGING;
    stagingDesc.BindFlags          = 0;
    stagingDesc.CPUAccessFlags     = D3D11_CPU_ACCESS_READ;
    stagingDesc.MiscFlags          = 0;
    
    for (int i = 0; i < 2; ++i)
    {
        HRESULT hr = m_pDevice->CreateTexture2D(&stagingDesc, nullptr, &m_pStagingTexture[i]);
        if (FAILED(hr))
        {
			$err, "Failed to create staging texture %1", i, hr;
            for (int j = 0; j < i; ++j)
            {
                if (m_pStagingTexture[j])
                {
                    m_pStagingTexture[j]->Release();
                    m_pStagingTexture[j] = nullptr;
                }
            }
            return false;
        }
    }
    
    // Cache dimensions
    m_lastWidth  = width;
    m_lastHeight = height;
    m_lastFormat = format;
    m_currentStagingIndex = 0;
    
    return true;
}



bool WindowCapture::createOrUpdateDIBSection(UINT width, UINT height)
{
    // Reuse existing DIB if dimensions match
    if (m_hDIBSection && m_dibWidth == width && m_dibHeight == height)
        return true;
    
    // Release old DIB
    if (m_hDIBSection)
    {
        DeleteObject(m_hDIBSection);
        m_hDIBSection = nullptr;
        m_pDIBBits = nullptr;
    }
    
    // Create new DIB Section
    BITMAPINFO bi              = {};
    bi.bmiHeader.biSize        = sizeof(BITMAPINFOHEADER);
    bi.bmiHeader.biWidth       = static_cast<LONG>(width);
    bi.bmiHeader.biHeight      = -static_cast<LONG>(height); // Top-down DIB
    bi.bmiHeader.biPlanes      = 1;
    bi.bmiHeader.biBitCount    = 32; // 32-bit BGRA
    bi.bmiHeader.biCompression = BI_RGB;
    bi.bmiHeader.biSizeImage   = 0;
    
    m_hDIBSection = CreateDIBSection(nullptr, &bi, DIB_RGB_COLORS, &m_pDIBBits, nullptr, 0);    
    if (!m_hDIBSection)
    {
		$err, "Failed to create DIB section", GetLastError();
        return false;
    }
    
    m_dibWidth  = width;
    m_dibHeight = height;
    
    return true;
}



bool WindowCapture::copyTextureDataToDIB(const D3D11_MAPPED_SUBRESOURCE& mapped, UINT width, UINT height)
{
    if (!m_pDIBBits)
        return false;
    
    const UINT8* pSrc = static_cast<const UINT8*>(mapped.pData);
    UINT8* pDst = static_cast<UINT8*>(m_pDIBBits);
    UINT srcPitch = mapped.RowPitch;
    UINT dstPitch = ((width * 32 + 31) / 32) * 4;
    
    UINT rowBytes = width * 4; // BGRA pixels
    
    // Add memory validation and alignment checks
    if (!pSrc || !pDst || rowBytes == 0 || height == 0)
        return false;    
    // Validate that we don't exceed buffer bounds
    if (rowBytes > dstPitch || rowBytes > srcPitch)
        return false;

    if (srcPitch == dstPitch)
    {
        // Single memcpy for entire image
        memcpy(pDst, pSrc, height * srcPitch);
        return true;
    }

    // Standard memcpy
    for (UINT y = 0; y < height; ++y)
    {
        const UINT8* srcRow = pSrc + y * srcPitch;
        UINT8* dstRow = pDst + y * dstPitch;
        
        // Validate pointers before memcpy
        if (srcRow && dstRow)
            memcpy(dstRow, srcRow, rowBytes);
    }
    
    return true;
}



bool WindowCapture::captureToPixelBuffer(void** ppPixels, UINT* pWidth, UINT* pHeight, UINT* pPitch)
{
    if (!IsWindow(m_hwnd))
        return false;

    HANDLE hDxSurface = nullptr;
    HRESULT hr = DwmGetDxSharedSurface(m_hwnd, &hDxSurface, NULL, NULL, NULL, NULL);
    if (FAILED(hr) || !hDxSurface)
    {
		$err, "Failed to get shared surface %1\nhDxSurface: %2", hr, hDxSurface;
        return false;
    }

    ID3D11Texture2D* pSharedTexture;
    hr = m_pDevice->OpenSharedResource(hDxSurface, IID_PPV_ARGS(&pSharedTexture));
    if (FAILED(hr) || !pSharedTexture)
    {
        $err, "Failed to open shared resource", hr;
		return false;
    }

    D3D11_TEXTURE2D_DESC desc;
    pSharedTexture->GetDesc(&desc);

    // Create or reuse staging textures and DIB
    if (!createOrUpdateStagingTextures(desc.Width, desc.Height, desc.Format) ||
        !createOrUpdateDIBSection(desc.Width, desc.Height))
    {
        pSharedTexture->Release();
        return false;
    }

    // **DOUBLE-BUFFER STRATEGY**: Use alternating staging textures
    ID3D11Texture2D* pCurrentStaging = m_pStagingTexture[m_currentStagingIndex];
    ID3D11Texture2D* pPreviousStaging = m_pStagingTexture[1 - m_currentStagingIndex];
    
    // Copy to current staging texture (GPU operation)
    m_pContext->CopyResource(pCurrentStaging, pSharedTexture);
    
    // Try to map the previous staging texture first
    // This avoids waiting for the current copy to complete
    D3D11_MAPPED_SUBRESOURCE mapped;
    ID3D11Texture2D* pTextureToCopy = pPreviousStaging;
    
    hr = m_pContext->Map(pPreviousStaging, 0, D3D11_MAP_READ, D3D11_MAP_FLAG_DO_NOT_WAIT, &mapped);
            
    if (hr == DXGI_ERROR_WAS_STILL_DRAWING)
    {
        // Previous texture still busy, try current texture with optimized retry                
        // Force completion of current copy
        m_pContext->Flush();
        
        // Try current staging texture with retry logic
        const int MAX_RETRIES = 3;
        int retryCount = 0;
        
        do 
        {
            hr = m_pContext->Map(pCurrentStaging, 0, D3D11_MAP_READ, D3D11_MAP_FLAG_DO_NOT_WAIT, &mapped);
            
            if (hr == DXGI_ERROR_WAS_STILL_DRAWING && retryCount < MAX_RETRIES)
            {
                retryCount++;
                if (retryCount == 1) // Quick yield
                    std::this_thread::yield();
                else if (retryCount == 2) // Short sleep
                    std::this_thread::sleep_for(std::chrono::microseconds(50));
                else
                {
                    // Blocking map as last resort
                    hr = m_pContext->Map(pCurrentStaging, 0, D3D11_MAP_READ, 0, &mapped);
                    break;
                }
            }
        } 
        while (hr == DXGI_ERROR_WAS_STILL_DRAWING && retryCount < MAX_RETRIES);
        
        pTextureToCopy = pCurrentStaging;
    }
     
	bool success = false;
    if (SUCCEEDED(hr))
    {
        if (copyTextureDataToDIB(mapped, desc.Width, desc.Height))
        {
            // Return direct pointer to pixel data
            *ppPixels = m_pDIBBits;
            *pWidth   = desc.Width;
            *pHeight  = desc.Height;
            *pPitch   = ((desc.Width * 32 + 31) / 32) * 4;                    
            success   = true;
        }
        else
            $err, "Failed to copy texture data to DIB section";
    
        // Unmap the texture we used
        m_pContext->Unmap(pTextureToCopy, 0);
    }
    else
		$err, "Failed to map staging texture", hr;
     
    // Swap staging buffer indices for next frame
    m_currentStagingIndex = 1 - m_currentStagingIndex; 
    pSharedTexture->Release();

	// Calling CloseHandle on the shared surface THROW EXCEPTION!!!
    //CloseHandle(hDxSurface);
    return success;
}



bool WindowCapture::captureToFile(const wchar_t* filename)
{
    void* pPixels = nullptr;
    UINT width, height, pitch;
    
    if (!captureToPixelBuffer(&pPixels, &width, &height, &pitch))
        return false;
    
    // Use existing DIB section with GDI+ for file saving
    Gdiplus::Bitmap bitmap(width, height, pitch, PixelFormat32bppARGB, static_cast<BYTE*>(pPixels));
    
    CLSID pngClsid;
    UINT num = 0, size = 0;
    Gdiplus::GetImageEncodersSize(&num, &size);
    if (size > 0)
    {
        auto pImageCodecInfo = std::make_unique<Gdiplus::ImageCodecInfo[]>(size / sizeof(Gdiplus::ImageCodecInfo));
        Gdiplus::GetImageEncoders(num, size, pImageCodecInfo.get());
        
        for (UINT j = 0; j < num; ++j)
        {
            if (wcscmp(pImageCodecInfo[j].MimeType, L"image/png") == 0)
            {
                pngClsid = pImageCodecInfo[j].Clsid;
                break;
            }
        }
    }
    
    Gdiplus::Status status = bitmap.Save(filename, &pngClsid, nullptr);
    return status == Gdiplus::Ok;
}



int main(int argc, char *argv[])
{
    SetPriorityClass(GetCurrentProcess(), HIGH_PRIORITY_CLASS);
    SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_TIME_CRITICAL);
    
    HWND hwnd = FindWindowA(NULL, "Device");    
    if (!hwnd)
    {
        $err, "Window not found";
        return 1;
    }

    WindowCapture capture;
    if (!capture.initialize(hwnd))
        return 1;

    std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
    std::chrono::milliseconds interval(1000);

    for (int i = 0; ; i++)
    {
        void* pPixels = nullptr;
        UINT width, height, pitch;
        
        if (capture.captureToPixelBuffer(&pPixels, &width, &height, &pitch))
        {
            wchar_t filename[256];
            swprintf_s(filename, L"captured\\capture_%06d.png", i);
            capture.captureToFile(filename);
        }
        else
			$err, "Capture failed (attempt %1)", i;
        
        startTime = startTime + interval;
        std::this_thread::sleep_until(startTime);
    }
    
    return 0;
}