filewriter_async.cpp #4

#define _CRT_SECURE_NO_DEPRECATE
#include <windows.h>

#include <vector>

#include "filewriter_async.h"

// Initial size of the write buffer per file reader, which will grow as needed
const int	cStartBufSize			= 64 * 1024;

// We need to keep track of the total allocated size of all write buffers, so we don't go over the maximum
static int	sAllocatedSize			= 0;
int FileWriterASync::mSectorSize	= 0;

FileWriterASync::BufInfo::BufInfo(FileWriterASync* inFileWriter) :
	mData(NULL),
	mLength(0),
	mSize(cStartBufSize)
{
	mData = (unsigned char*) ::VirtualAlloc(NULL, inFileWriter->mMaxMemSize / 2, MEM_RESERVE, PAGE_READWRITE);
	::VirtualAlloc(mData, mSize, MEM_COMMIT, PAGE_READWRITE);
	sAllocatedSize += mSize;
	::memset(&mOverlappedInfo, 0, sizeof(OVERLAPPED));
	mOverlappedInfo.hEvent = ::CreateEvent(NULL, TRUE, FALSE, NULL);
}


FileWriterASync::BufInfo::~BufInfo()
{
	sAllocatedSize -= mSize;
	::VirtualFree(mData, 0, MEM_RELEASE);
	::CloseHandle(mOverlappedInfo.hEvent);
}


FileWriterASync* FileWriterASync::GetFileWriter(const std::string& inName, long long inFileSize)
{
	if (mSectorSize == 0)
	{
		std::string root_path(inName, 0, 3);
		DWORD sector_size;
		::GetDiskFreeSpace(root_path.c_str(), NULL, &sector_size, NULL, NULL);
		mSectorSize = sector_size;
		if (mSectorSize == 0)
			mSectorSize = 4096; // Not sure what went wrong, but this should be a safe default
	}	
	return new FileWriterASync(inName, inFileSize);
}


void FileWriterASync::Finalize()
{
}


FileWriterASync::FileWriterASync(const std::string& inName, long long inFileSize) :
	FileWriter(inName, inFileSize),
	mHandle(INVALID_HANDLE_VALUE),
	mWriteOffset(0)
{
}


FileWriterASync::~FileWriterASync()
{
}


bool FileWriterASync::Open(bool inWriteable, int inModTime)
{
	mWriteOffset = 0;
	BufInfo* buf_info = new BufInfo(this);
	mBuffers.push_back(buf_info);
	mHandle = CreateFile(inWriteable, true, inModTime);
	return mHandle != INVALID_HANDLE_VALUE;
}


bool FileWriterASync::Close()
{
	return Flush(true) >= 0;
}


void FileWriterASync::GetWriteBuffer(unsigned char*& outBuf, int& outLength)
{
	BufInfo* buf_info = mBuffers.back();
	outBuf = buf_info->mData + buf_info->mLength;
	outLength = buf_info->mSize - buf_info->mLength;
}


bool FileWriterASync::ReleaseWriteBuffer(int inRemainingLength)
{
	BufInfo* cur_buf_info = mBuffers.back();
	cur_buf_info->mLength = cur_buf_info->mSize - inRemainingLength;
	// We need to have at least two bytes free, as we may need to do a line-end conversion
	if (inRemainingLength <= 1)
	{
		if (cur_buf_info->mSize >= mMaxMemSize / 2)
		{
			unsigned char* left_bytes = 0;
			int bytes_left = Flush(false);
			if (bytes_left > 0)
			{
				left_bytes = new unsigned char[bytes_left];
				::memcpy(left_bytes, cur_buf_info->mData + cur_buf_info->mLength - bytes_left, bytes_left);
			}
			while (sAllocatedSize > mMaxMemSize / 2)
				Wait(1);
			BufInfo* new_buf_info = new BufInfo(this);
			mBuffers.push_back(new_buf_info);
			if (bytes_left > 0)
			{
				::memcpy(new_buf_info->mData, left_bytes, bytes_left);
				delete[] left_bytes;
				new_buf_info->mLength = bytes_left;
			}
			return bytes_left >= 0;
		}
		else
		{
			BufInfo* new_buf_info = mBuffers.back();
			new_buf_info->mSize *= 2;
			::VirtualAlloc(new_buf_info->mData, new_buf_info->mSize, MEM_COMMIT, PAGE_READWRITE);
			sAllocatedSize += new_buf_info->mSize / 2;
		}
	}
	return true;
}


int FileWriterASync::Flush(bool inWriteAll)
{
	Log("Flush\n");
	BufInfo* buf_info = mBuffers.back();	
	int write_length = 0;
	if (buf_info->mLength > 0)
	{
		buf_info->mOverlappedInfo.Offset = (DWORD)(mWriteOffset);
		buf_info->mOverlappedInfo.OffsetHigh = (DWORD)(mWriteOffset >> 32);
		if (inWriteAll)
			write_length = ((buf_info->mLength + mSectorSize - 1) / mSectorSize) * mSectorSize;
		else
			write_length = ((buf_info->mLength - mSectorSize + 1) / mSectorSize) * mSectorSize;
		if (::WriteFile(mHandle, buf_info->mData, write_length, 0, &buf_info->mOverlappedInfo) == 0)
		{
			if (GetLastError() != ERROR_IO_PENDING)
				return -1;
		}
		int bytes_left = buf_info->mLength - write_length;
		if (bytes_left < 0)
		{
			bytes_left = 0; // This must have been the final flush call
			mWriteOffset += buf_info->mLength;
		}
		else
			mWriteOffset += write_length;
		return bytes_left;
	}
	else
	{
		delete buf_info;
		mBuffers.pop_back();
		return 0;
	}
}


void FileWriterASync::Wait(int inMSecs)
{
	if (!mBuffers.empty())
	{
		std::vector<HANDLE> handle_vec;
		for (std::list<BufInfo*>::iterator it = mBuffers.begin(); it != mBuffers.end(); ++it)
			if ((*it)->mLength > 0)
				handle_vec.push_back((*it)->mOverlappedInfo.hEvent);
		if (handle_vec.size() > 0)
			::WaitForMultipleObjects((DWORD)handle_vec.size(), &handle_vec[0], FALSE, inMSecs);
		for (std::list<BufInfo*>::iterator it = mBuffers.begin(); it != mBuffers.end(); )
		{
			bool remove = true;
			if ((*it)->mLength > 0)
			{
				DWORD bytes_written;
				remove = (GetOverlappedResult(mHandle, &(*it)->mOverlappedInfo, &bytes_written, FALSE) == TRUE);
			}
			if (remove)
			{
				std::list<BufInfo*>::iterator next = it;
				++next;
				delete *it;
				mBuffers.erase(it);
				it = next;
			}
			else
			{
				++it;
			}
		}
	}
}


void FileWriterASync::Finish()
{
	while (!mBuffers.empty())
		Wait(INFINITE);
	CloseFile(mHandle, mWriteOffset);
}


bool FileWriterASync::IsFinished()
{
	Wait(1);
	if (mBuffers.empty())
	{
		CloseFile(mHandle, mWriteOffset);
		return true;
	}
	return false;
}


void FileWriterASync::HandleAbort()
{
	::CancelIo(mHandle);
	for (std::list<BufInfo*>::iterator it = mBuffers.begin(); it != mBuffers.end(); ++it)
		delete *it;
	mBuffers.clear();
	CloseFile(mHandle);
	// Aborted before file actions could be finished, so delete
	DeleteFile();
}


bool FileWriterASync::ChModTime(int inModTime)
{
	return FileWriter::ChModTime(mHandle, inModTime);
}