#include <iprt/file.h>

#include <iprt/err.h>

#include <iprt/mem.h>

#include <iprt/param.h>

#include <iprt/string.h>

#include <iprt/test.h>

#define TSTFILEAIO_MAX_REQS_IN_FLIGHT 64

#define TSTFILEAIO_BUFFER_SIZE (64*_1K)

static RTTEST g_hTest = NIL_RTTEST;

void tstFileAioTestReadWriteBasic(RTFILE File, bool fWrite, void *pvTestBuf,

size_t cbTestBuf, size_t cbTestFile, uint32_t cMaxReqsInFlight)

{

/* Allocate request array. */

RTFILEAIOREQ *paReqs;

paReqs = (PRTFILEAIOREQ)RTTestGuardedAllocHead(g_hTest, cMaxReqsInFlight * sizeof(RTFILEAIOREQ));

RTTESTI_CHECK_RETV(paReqs);

RT_BZERO(paReqs, sizeof(cMaxReqsInFlight * sizeof(RTFILEAIOREQ)));

/* Allocate array holding pointer to data buffers. */

void **papvBuf = (void **)RTTestGuardedAllocHead(g_hTest, cMaxReqsInFlight * sizeof(void *));

RTTESTI_CHECK_RETV(papvBuf);

/* Allocate the buffers*/

for (unsigned i = 0; i < cMaxReqsInFlight; i++)

{

RTTESTI_CHECK_RC_OK_RETV(RTTestGuardedAlloc(g_hTest, cbTestBuf, PAGE_SIZE, true /*fHead*/, &papvBuf[i]));

if (fWrite)

memcpy(papvBuf[i], pvTestBuf, cbTestBuf);

if (fWrite)

memcpy(papvBuf[i], pvTestBuf, cbTestBuf);

else

RT_BZERO(papvBuf[i], cbTestBuf);

}

/* Allocate array holding completed requests. */

RTFILEAIOREQ *paReqsCompleted;

paReqsCompleted = (PRTFILEAIOREQ)RTTestGuardedAllocHead(g_hTest, cMaxReqsInFlight * sizeof(RTFILEAIOREQ));

RTTESTI_CHECK_RETV(paReqsCompleted);

RT_BZERO(paReqsCompleted, cMaxReqsInFlight * sizeof(RTFILEAIOREQ));

/* Create a context and associate the file handle with it. */

RTFILEAIOCTX hAioContext;

RTTESTI_CHECK_RC_RETV(RTFileAioCtxCreate(&hAioContext, cMaxReqsInFlight), VINF_SUCCESS);

RTTESTI_CHECK_RC_RETV(RTFileAioCtxAssociateWithFile(hAioContext, File), VINF_SUCCESS);

/* Initialize requests. */

for (unsigned i = 0; i < cMaxReqsInFlight; i++)

RTFileAioReqCreate(&paReqs[i]);

RTFOFF off = 0;

int cRuns = 0;

uint64_t NanoTS = RTTimeNanoTS();

size_t cbLeft = cbTestFile;

while (cbLeft)

{

int rc;

int cReqs = 0;

for (unsigned i = 0; i < cMaxReqsInFlight; i++)

{

size_t cbTransfer = cbLeft < cbTestBuf ? cbLeft : cbTestBuf;

if (!cbTransfer)

break;

if (fWrite)

rc = RTFileAioReqPrepareWrite(paReqs[i], File, off, papvBuf[i],

cbTransfer, papvBuf[i]);

else

rc = RTFileAioReqPrepareRead(paReqs[i], File, off, papvBuf[i],

cbTransfer, papvBuf[i]);

RTTESTI_CHECK_RC(rc, VINF_SUCCESS);

cbLeft -= cbTransfer;

off += cbTransfer;

cReqs++;

}

rc = RTFileAioCtxSubmit(hAioContext, paReqs, cReqs);

RTTESTI_CHECK_MSG(rc == VINF_SUCCESS, ("Failed to submit tasks after %d runs. rc=%Rrc\n", cRuns, rc));

if (rc != VINF_SUCCESS)

break;

/* Wait */

uint32_t cCompleted = 0;

RTTESTI_CHECK_RC(rc = RTFileAioCtxWait(hAioContext, cReqs, RT_INDEFINITE_WAIT,

paReqsCompleted, cMaxReqsInFlight, &cCompleted),

VINF_SUCCESS);

if (rc != VINF_SUCCESS)

break;

if (!fWrite)

{

for (uint32_t i = 0; i < cCompleted; i++)

{

/* Compare that we read the right stuff. */

void *pvBuf = RTFileAioReqGetUser(paReqsCompleted[i]);

RTTESTI_CHECK(pvBuf);

size_t cbTransfered;

RTTESTI_CHECK_RC(rc = RTFileAioReqGetRC(paReqsCompleted[i], &cbTransfered), VINF_SUCCESS);

if (rc != VINF_SUCCESS)

break;

RTTESTI_CHECK_MSG(cbTransfered == cbTestBuf, ("cbTransfered=%zd\n", cbTransfered));

RTTESTI_CHECK_RC_OK(rc = (memcmp(pvBuf, pvTestBuf, cbTestBuf) == 0 ? VINF_SUCCESS : VERR_BAD_EXE_FORMAT));

if (rc != VINF_SUCCESS)

break;

memset(pvBuf, 0, cbTestBuf);

}

}

cRuns++;

if (RT_FAILURE(rc))

break;

}

NanoTS = RTTimeNanoTS() - NanoTS;

uint64_t SpeedKBs = (uint64_t)(cbTestFile / (NanoTS / 1000000000.0) / 1024);

RTTestValue(g_hTest, "Throughput", SpeedKBs, RTTESTUNIT_KILOBYTES_PER_SEC);

/* cleanup */

for (unsigned i = 0; i < cMaxReqsInFlight; i++)

RTTestGuardedFree(g_hTest, papvBuf[i]);

RTTestGuardedFree(g_hTest, papvBuf);

for (unsigned i = 0; i < cMaxReqsInFlight; i++)

RTTESTI_CHECK_RC(RTFileAioReqDestroy(paReqs[i]), VINF_SUCCESS);

RTTESTI_CHECK_RC(RTFileAioCtxDestroy(hAioContext), VINF_SUCCESS);

RTTestGuardedFree(g_hTest, paReqs);

}

int main()

{

int rc = RTTestInitAndCreate("tstRTFileAio", &g_hTest);

if (rc)

return rc;

/* Check if the API is available. */

RTTestSub(g_hTest, "RTFileAioGetLimits");

RTFILEAIOLIMITS AioLimits;

RT_ZERO(AioLimits);

RTTESTI_CHECK_RC(rc = RTFileAioGetLimits(&AioLimits), VINF_SUCCESS);

if (RT_SUCCESS(rc))

{

RTTestSub(g_hTest, "Write");

RTFILE hFile;

RTTESTI_CHECK_RC(rc = RTFileOpen(&hFile, "tstFileAio#1.tst",

RTFILE_O_READWRITE | RTFILE_O_CREATE_REPLACE | RTFILE_O_DENY_NONE | RTFILE_O_ASYNC_IO),

VINF_SUCCESS);

if (RT_SUCCESS(rc))

{

uint8_t *pbTestBuf = (uint8_t *)RTTestGuardedAllocTail(g_hTest, TSTFILEAIO_BUFFER_SIZE);

for (unsigned i = 0; i < TSTFILEAIO_BUFFER_SIZE; i++)

pbTestBuf[i] = i % 256;

uint32_t cReqsMax = AioLimits.cReqsOutstandingMax < TSTFILEAIO_MAX_REQS_IN_FLIGHT

? AioLimits.cReqsOutstandingMax

: TSTFILEAIO_MAX_REQS_IN_FLIGHT;

/* Basic write test. */

RTTestIPrintf(RTTESTLVL_ALWAYS, "Preparing test file, this can take some time and needs quite a bit of harddisk space...\n");

tstFileAioTestReadWriteBasic(hFile, true /*fWrite*/, pbTestBuf, TSTFILEAIO_BUFFER_SIZE, 100*_1M, cReqsMax);

/* Reopen the file before doing the next test. */

RTTESTI_CHECK_RC(RTFileClose(hFile), VINF_SUCCESS);

if (RTTestErrorCount(g_hTest) == 0)

{

RTTestSub(g_hTest, "Read/Write");

RTTESTI_CHECK_RC(rc = RTFileOpen(&hFile, "tstFileAio#1.tst",

RTFILE_O_READWRITE | RTFILE_O_OPEN | RTFILE_O_DENY_NONE | RTFILE_O_ASYNC_IO),

VINF_SUCCESS);

if (RT_SUCCESS(rc))

{

tstFileAioTestReadWriteBasic(hFile, false /*fWrite*/, pbTestBuf, TSTFILEAIO_BUFFER_SIZE, 100*_1M, cReqsMax);

RTFileClose(hFile);

}

}

/* Cleanup */

RTFileDelete("tstFileAio#1.tst");

}

}

/*

return RTTestSummaryAndDestroy(g_hTest);

}