|| || || || || ||

---

IOBuffer.h

Go to the documentation of this file.

#ifndef Impala_Util_IOBuffer_h
#define Impala_Util_IOBuffer_h

#include <stdio.h>
#include <string>
#include <deque>
#include <iostream>
#include "Basis/ILog.h"
#include "Util/Channel.h"
#include "Basis/NativeTypeFormats.h"
#include "Link/Mpi/MpiFuncs.h"
#ifdef BZIP2_USED
#include "bzlib.h"
#endif // BZIP2_USED
#ifdef ZLIB_USED
#ifdef WIN32
#include "Png/zlib.h"
#else
#include "zlib.h"
#endif
#endif // ZLIB_USED


namespace Impala
{
namespace Util
{

class IOBuffer
{
public:
    typedef Int64 PositionType;

    static const UInt32 cBzipBlockSize = 16 * 1024 * 1024;

    IOBuffer()
    {
        mBuffer = 0;
        mSize = 0;
        mPosition = 0;
        mDataChannel = 0;
    }
 
    IOBuffer(PositionType length)
    {
        mBuffer = new unsigned char[length];
        mSize = length;
        mPosition = 0;
        mDataChannel = 0;
    }
 
    IOBuffer(PositionType length, unsigned char* buffer)
    {
        mBuffer = buffer;
        mSize = length;
        mPosition = 0;
        mDataChannel = 0;
    }

    IOBuffer(String s)
    {
        mSize = s.size();
        mBuffer = new unsigned char [mSize+1];
        std::memcpy(mBuffer, s.c_str(), mSize);
        mBuffer[mSize] = 0;
        mPosition = 0;
        mDataChannel = 0;
    }

    virtual
    ~IOBuffer()
    {
        if (mDataChannel)
            WriteIOBufferToChannel();
        if (mBuffer)
            delete [] mBuffer;
        if (!mDataChannelTmp.empty())
            remove(mDataChannelTmp.c_str());
    }

    virtual bool
    Valid()
    {
        return true;
    }

    virtual PositionType
    Available()
    {
        PositionType res = mSize - mPosition;
        if (res < 0)
        {
            ILOG_ERROR("Available: position beyond size");
            return 0;
        }
        if (res == 1)
            return 0;
        return res;
    }

    virtual int
    AvailableInt32()
    {
        return Min<Int64>(Available(), 0x7fffffff);
    }

    virtual PositionType
    Size()
    {
        return mSize;
    }

    virtual void
    SetSize(PositionType size)
    {
        mSize = size;
    }

    template<class BackInsertIterator>
    void
    ReadStrings(BackInsertIterator bi, bool skipComments, int count)
    {
        if (! Valid())
            return;

        int c=0;
        while (Available())
        {
            String line = ReadLine();
            if (line[0] || !skipComments)
            {
                if ((line[0] != '#') || !skipComments)
                    *bi++ = line;
            }
            if (++c==count)
                break;
        }
    }

    template <class Iterator>
    bool
    WriteStrings(Iterator begin, Iterator end, int count = -1)
    {
        if (! Valid())
            return false;

        int c=0;
        for (Iterator it=begin ; it!=end ; it++)
        {
            String s = *it;
            Puts(s);
            if (++c == count)
                break;
        }
        return true;
    }


    virtual Int64
    Read(void* buf, Int64 bytesToRead)
    {
        Int64 available = Available();
        if (available < bytesToRead)
            bytesToRead = available;
        memcpy(buf, mBuffer + mPosition, bytesToRead);
        mPosition += bytesToRead;
        return bytesToRead;
    }

#ifdef BZIP2_USED
    Int64
    ReadBzip(void* buf, Int64 bytesToRead)
    {
        UInt32 nrBlocks = 0;
        Int64 nrBytes = 0;
        Read(&nrBlocks, sizeof(UInt32));
        Read(&nrBytes, sizeof(Int64));
        if (nrBytes > bytesToRead)
        {
            ILOG_ERROR("ReadBzip: Buffer to small");
            return 0;
        }

        char* bufPtr = (char *) buf;
        char* bzipBuf = new char[cBzipBlockSize];
        Int64 left = nrBytes;
        Int64 nrRead = 0;
        for (int i=0 ; i<nrBlocks ; i++)
        {
            UInt32 bzipSize = 0;
            Read(&bzipSize, sizeof(UInt32));
            Read(bzipBuf, bzipSize);

            UInt32 nr = (left > cBzipBlockSize) ? cBzipBlockSize : left;
            int res = BZ2_bzBuffToBuffDecompress(bufPtr, &nr, bzipBuf,
                                                 bzipSize, 0, 0);
            if (res == BZ_OK)
                nrRead += nr;
            else
                ILOG_ERROR("ReadBzip: Decompression failed");
            bufPtr += nr;
            left -= nr;
        }
        delete bzipBuf;
        return nrRead;
    }
#else
    Int64
    ReadBzip(void* buf, Int64 bytesToRead)
    {
        ILOG_ERROR("ReadBzip not supported");
        return 0;
    }
#endif // BZIP2_USED

#ifdef ZLIB_USED
    Int64
    ReadZlib(void* buf, Int64 bytesToRead)
    {
        UInt32 nrBlocks = 0;
        Int64 nrBytes = 0;
        Read(&nrBlocks, sizeof(UInt32));
        Read(&nrBytes, sizeof(Int64));
        if (nrBytes > bytesToRead)
        {
            ILOG_ERROR("ReadZlib: Buffer to small");
            return 0;
        }

        unsigned char* bufPtr = (unsigned char *) buf;
        unsigned char* bzipBuf = new unsigned char[cBzipBlockSize];
        Int64 left = nrBytes;
        Int64 nrRead = 0;
        for (int i=0 ; i<nrBlocks ; i++)
        {
            UInt32 bzipSize = 0;
            Read(&bzipSize, sizeof(UInt32));
            Read(bzipBuf, bzipSize);

            uLongf nr = (left > cBzipBlockSize) ? cBzipBlockSize : left;
            int res = uncompress(bufPtr, &nr, bzipBuf, bzipSize);
            if (res == Z_OK)
                nrRead += nr;
            else
                ILOG_ERROR("ReadZlib: Decompression failed");
            bufPtr += nr;
            left -= nr;
        }
        delete bzipBuf;
        return nrRead;
    }
#else
    Int64
    ReadZlib(void* buf, Int64 bytesToRead)
    {
        ILOG_ERROR("ReadZlib not supported");
        return 0;
    }
#endif // ZLIB_USED

    virtual String
    ReadLine()
    {
        Int64 start = GetPosition();
        Int64 i = 0;
        Int64 available = Available();
        while ((i < available) && (mBuffer[start+i] != '\n'))
            i++;
        Int64 skipN = (mBuffer[start+i] == '\n') ? 1 : 0;
        SetPosition(start + i + skipN);
        if ((i > 0) && (mBuffer[start+i-1] == '\r')) // strip '\r' if present
            i--;
        const char* p = reinterpret_cast<const char*>(mBuffer + start);
        return String(p, i);
    }

    virtual Int64
    Gets(char* buf, Int64 bytesToRead)
    {
        Int64 available = Available();
        if (available < bytesToRead)
            bytesToRead = available;
        PositionType start = mPosition;
        Int64 i = 0;
        while ((i < bytesToRead) && (mBuffer[start+i] != '\n'))
        {
            buf[i] = mBuffer[start+i];
            i++;
        }
        bytesToRead = i;
        mPosition += bytesToRead;
        if (mBuffer[start+i] == '\n')
            mPosition++;
        return bytesToRead;
    }

    virtual void
    NativeTypeRead(Int8* ptr)
    {
        DoNativeTypeRead(ptr);
    }

    virtual void
    NativeTypeRead(UInt8* ptr)
    {
        DoNativeTypeRead(ptr);
    }

    virtual void
    NativeTypeRead(Int16* ptr)
    {
        DoNativeTypeRead(ptr);
    }

    virtual void
    NativeTypeRead(UInt16* ptr)
    {
        DoNativeTypeRead(ptr);
    }

    virtual void
    NativeTypeRead(Int32* ptr)
    {
        DoNativeTypeRead(ptr);
    }

    virtual void
    NativeTypeRead(UInt32* ptr)
    {
        DoNativeTypeRead(ptr);
    }
    
    virtual void
    NativeTypeRead(Int64* ptr)
    {
        DoNativeTypeRead(ptr);
    }

    virtual void
    NativeTypeRead(UInt64* ptr)
    {
        DoNativeTypeRead(ptr);
    }

    virtual void
    NativeTypeRead(Real32* ptr)
    {
        DoNativeTypeRead(ptr);
    }

    virtual void
    NativeTypeRead(Real64* ptr)
    {
        DoNativeTypeRead(ptr);
    }

    virtual Int64
    Write(const void* buf, Int64 bytesToWrite)
    {
        Int64 available = Available();
        if (available < bytesToWrite)
            bytesToWrite = available;
        memcpy(GetBuffer() + GetPosition(), buf, bytesToWrite);
        SetPosition(GetPosition() + bytesToWrite);
        return bytesToWrite;
    }

#ifdef BZIP2_USED
    void
    WriteBzip(void* srcBuf, Int64 bytesToWrite)
    {
        UInt32 nrBlocks = bytesToWrite / cBzipBlockSize;
        if (bytesToWrite % cBzipBlockSize != 0)
            nrBlocks++;
        Write(&nrBlocks, sizeof(UInt32));
        Write(&bytesToWrite, sizeof(Int64));

        // BZ man : To guarantee that the compressed data will fit in its buffer,
        // allocate an output buffer of size 1% larger than the uncompressed
        // data, plus six hundred extra bytes.
        const UInt32 cBzipBufCap = cBzipBlockSize * 1.01 + 600;
        char* srcBufPtr = (char *) srcBuf;
        char* bzipBuf = new char[cBzipBufCap];
        Int64 left = bytesToWrite;
        for (int i=0 ; i<nrBlocks ; i++)
        {
            Int64 nr = left;
            if (nr > cBzipBlockSize)
                nr = cBzipBlockSize;
            UInt32 bzipSize = cBzipBufCap;
            int res = BZ2_bzBuffToBuffCompress(bzipBuf, &bzipSize, srcBufPtr,
                                               nr, 9, 1, 30);
            if (res == BZ_OK)
            {
                Write(&bzipSize, sizeof(UInt32));
                Write(bzipBuf, bzipSize);
            }
            else
            {
                ILOG_ERROR("WriteBzip: Compression failed");
            }
            srcBufPtr += nr;
            left -= nr;
        }
        delete bzipBuf;
    }
#else
    void
    WriteBzip(void* srcBuf, Int64 bytesToWrite)
    {
        ILOG_ERROR("WriteBzip not supported");
    }
#endif // BZIP2_USED

#ifdef ZLIB_USED
    void
    WriteZlib(void* srcBuf, Int64 bytesToWrite)
    {
        UInt32 nrBlocks = bytesToWrite / cBzipBlockSize;
        if (bytesToWrite % cBzipBlockSize != 0)
            nrBlocks++;
        Write(&nrBlocks, sizeof(UInt32));
        Write(&bytesToWrite, sizeof(Int64));

        // BZ man : To guarantee that the compressed data will fit in its buffer,
        // allocate an output buffer of size 1% larger than the uncompressed
        // data, plus six hundred extra bytes.
        //const UInt32 cBzipBufCap = cBzipBlockSize * 1.01 + 600;
        const UInt32 cBzipBufCap = compressBound(cBzipBlockSize) + 32;
        unsigned char* srcBufPtr = (unsigned char *) srcBuf;
        unsigned char* bzipBuf = new unsigned char[cBzipBufCap];
        Int64 left = bytesToWrite;
        for (int i=0 ; i<nrBlocks ; i++)
        {
            Int64 nr = left;
            if (nr > cBzipBlockSize)
                nr = cBzipBlockSize;
            uLongf bzipSize = cBzipBufCap;
            int res = compress(bzipBuf, &bzipSize, srcBufPtr, nr);
            if (res == Z_OK)
            {
                Write(&bzipSize, sizeof(UInt32));
                Write(bzipBuf, bzipSize);
            }
            else
            {
                ILOG_ERROR("WriteZlib: Compression failed");
            }
            srcBufPtr += nr;
            left -= nr;
        }
        delete bzipBuf;
    }
#else
    void
    WriteZlib(void* srcBuf, Int64 bytesToWrite)
    {
        ILOG_ERROR("WriteZlib not supported");
    }
#endif // BZIP2_USED

/*
    virtual void
    WriteLine(const String& line)
    {
        ILOG_ERROR("Not implemented yet, involves reallocing..."
                   << "this is duplication of available code");
    }
*/

    virtual Int64
    Puts(const char* buf)
    {
        Int64 len = strlen(buf);
        Int64 available = Available();
        if (available < len)
            len = available;
        unsigned char* start = GetBuffer() + GetPosition();
        memcpy(start, buf, len);
        start[len] = '\n';
        SetPosition(GetPosition() + len + 1);
        return 1;
    }

    virtual Int64
    Puts(const String& str)
    {
        return Puts(str.c_str());
    }

    virtual void
    NativeTypeWrite(Int8 val)
    {
        DoNativeTypeWrite(val);
    }

    virtual void
    NativeTypeWrite(UInt8 val)
    {
        DoNativeTypeWrite(val);
    }

    virtual void
    NativeTypeWrite(Int16 val)
    {
        DoNativeTypeWrite(val);
    }

    virtual void
    NativeTypeWrite(UInt16 val)
    {
        DoNativeTypeWrite(val);
    }

    virtual void
    NativeTypeWrite(Int32 val)
    {
        DoNativeTypeWrite(val);
    }

    virtual void
    NativeTypeWrite(UInt32 val)
    {
        DoNativeTypeWrite(val);
    }
    
    virtual void
    NativeTypeWrite(Int64 val)
    {
        DoNativeTypeWrite(val);
    }

    virtual void
    NativeTypeWrite(UInt64 val)
    {
        DoNativeTypeWrite(val);
    }

    virtual void
    NativeTypeWrite(Real32 val)
    {
        DoNativeTypeWrite(val);
    }

    virtual void
    NativeTypeWrite(Real64 val)
    {
        DoNativeTypeWrite(val);
    }



    virtual void
    SetPosition(PositionType position)
    {
        mPosition = position;
    }

    // For specializations that allow for a growing buffer
    virtual void
    SetPositionAndSize(PositionType position)
    {
        mPosition = position;
        if (mPosition > mSize)
            mSize = position;
    }

    virtual PositionType
    GetPosition()
    {
        return mPosition;
    }

    virtual void
    Rewind()
    {
        mPosition = 0;
    }

    virtual void
    FastForward()
    {
        mPosition = mSize;
    }

    virtual PositionType
    Seek(PositionType offset, int whence)
    {
        if (whence == SEEK_CUR)
            SetPosition(GetPosition() + offset);
        else if (whence == SEEK_END)
            SetPosition(Size() + offset);
        else if (whence == SEEK_SET)
            SetPosition(offset);
        else
            return -1;

        return GetPosition();
    }


    unsigned char*
    GetBuffer()
    {
        return mBuffer;
    }

    void
    SetBuffer(unsigned char* buf, PositionType size)
    {
        mBuffer = buf;
        mSize = size;
    }

    void
    SetUseChannel(Channel* channel, String channelFile, String tmpFile)
    {
        mDataChannel = channel;
        mDataChannelFile = channelFile;
        mDataChannelTmp = tmpFile;
    }

protected:

    template <class NativeType>
    void
    DoNativeTypeRead(NativeType* ptr)
    {
        String fs = NativeTypeFormat<NativeType>(0);
        /* This seems to copy all of the remaining buffer :-(
        const char* start = (const char*) (GetBuffer() + GetPosition());
        sscanf(start, fs.c_str(), ptr);
        int i = 0;
        while (start[i] != ' ')
            i++;
        */
        char buf[1024];
        unsigned char* src = GetBuffer() + GetPosition();
        Int64 i = 0;
        while (src[i] != ' ')
            buf[i++] = src[i];
        buf[i] = ' ';
        sscanf(buf, fs.c_str(), ptr);
        SetPosition(GetPosition() + i + 1);
    }

    template <class NativeType>
    void
    DoNativeTypeWrite(NativeType val)
    {
        String fs = NativeTypeFormat<NativeType>(0) + " ";
        char buf[1024];
        sprintf(buf, fs.c_str(), val);
        Int64 i = 0;
        while (buf[i] != ' ')
            i++;
        Write(buf, i+1);
    }

    bool
    WriteIOBufferToChannel()
    {
        SetSize(GetPosition()); // assume the last write designates the end
        char* buf = mDataChannel->Buffer();
        sprintf(buf, "openfilebuffer \"%s\" 0\0", mDataChannelFile.c_str());
        mDataChannel->SendRequest(strlen(buf)+1);
        if (mDataChannel->LastSendHadError())
        {
            ILOG_ERROR("WriteIOBufferToChannel: open failed for [" <<
                       mDataChannelFile << "]");
            return false;
        }
        int id;
        Int64 bufSize;
        sscanf(buf, "%d %lld", &id, &bufSize); // bufSize is dummy when writing
        ILOG_DEBUG(mDataChannelFile << ", id = " << id);

        Int64 nrWritten = 0;
        SetPosition(0);
        while (Available() > 0)
        {
            sprintf(buf, "writefilebuffer %d\0", id);
            int used = strlen(buf) + 1;
            int nrData = Min<Int64>(Channel::DATA_BUFFER_SIZE-used, Available());
            Read(buf+used, nrData);
            int len = mDataChannel->SendRequest(used+nrData);
            if (mDataChannel->LastSendHadError())
            {
                ILOG_ERROR("WriteIOBufferToChannel: write failed for [" <<
                           mDataChannelFile << "]");
                break;
            }
            nrWritten += nrData;
            ILOG_DEBUG("  wrote " << nrData << " bytes, total = " << nrWritten);
        }
        sprintf(buf, "closefilebuffer %d\0", id);
        mDataChannel->SendRequest(strlen(buf)+1);
        return true;
    }

    void
    CheckDataChannelWrite()
    {
        if (mDataChannel != 0)
        {
            WriteIOBufferToChannel();
            mDataChannel = 0;
        }
    }

private:

    unsigned char* mBuffer;
    PositionType   mSize;
    PositionType   mPosition;

    Channel*       mDataChannel;
    String         mDataChannelFile; // "logical" name of the file at the server
    String         mDataChannelTmp; // name of the temporary file

    ILOG_VAR_DEC;
};

ILOG_VAR_INIT(IOBuffer, Impala.Util);

inline IOBuffer*
CreateIOBufferFromFile(CString filename)
{
    // read all contents of file and put in the IOBuffer
    FILE* file = fopen(filename.c_str(), "r");
    if (!file)
    {
        std::cout << "[CreateIOBufferFromFile] failed to open " << filename
                  << std::endl;
        return 0;
    }

    int blockSize = 0x1000;
    std::deque<unsigned char*> bufferList;
    int read;
    do
    {
        unsigned char* buffer = new unsigned char[blockSize];
        read = fread(buffer,1,blockSize,file);
        bufferList.push_back(buffer);
    }
    while (read == blockSize);
    fclose(file);

    // create total buffer
    int total = ((bufferList.size()-1)*blockSize) + read;
    unsigned char* buffer = new unsigned char[total];
    int i;
    for(i=0 ; i<bufferList.size() - 1 ; i++)
    {
        memcpy(buffer + (i*blockSize), bufferList[i], blockSize);
        delete [] bufferList[i];
        bufferList[i] = 0;
    }
    memcpy(buffer + (i*blockSize), bufferList[i], read); // last block only #read bytes
    delete [] bufferList[i];
    bufferList[i] = 0;

    return new IOBuffer(total, buffer);
}

inline void
Broadcast(IOBuffer* buffer, int root)
{
#ifdef MPI_USED
    ILOG_VAR(Impala.Util.Broadcast);
    MPI_Barrier(MPI_COMM_WORLD);
    Int64 size = buffer->Size();
    MPI_Bcast(&size, 1, MPI_LONG_LONG, root, MPI_COMM_WORLD);
    unsigned char* sendBuffer;
    if (Link::Mpi::MyId() == root)
        sendBuffer = buffer->GetBuffer();
    else
        sendBuffer = new unsigned char[size];
    ILOG_DEBUG(Link::Mpi::MyId() << ": about to broadcast, buf = " <<
               (void*)sendBuffer << "size = " << size);
    unsigned char* curPos = sendBuffer;
    Int64 transmitted = 0;
    while (transmitted < size)
    {
        int maxBlock = 1024 * 1024 * 1024;
        int blockSize = (size - transmitted > maxBlock) ? maxBlock 
                                                        : size - transmitted;
        MPI_Bcast(curPos, blockSize, MPI_CHAR, root, MPI_COMM_WORLD);
        curPos += blockSize;
        transmitted += blockSize;
    }
    if (Link::Mpi::MyId() != root)
    {
        buffer->SetBuffer(sendBuffer, size);
    }
    ILOG_DEBUG("done");
#endif // MPI_USED
}

} // namespace Util
} // namespace Impala

#endif // Impala_Util_IOBuffer_h

---