muduo库阅读（37）——Net部分：压缩数据流ZlibOutputStream

/*
* 压缩数据流
* 输入是普通数据流，输出是压缩数据流
*/
class ZlibOutputStream : boost::noncopyable
{
public:
// 初始化
explicit ZlibOutputStream(Buffer* output)
: output_(output),
zerror_(Z_OK),
bufferSize_(1024)
{
bzero(&zstream_, sizeof zstream_);

// 初始化压缩数据流
zerror_ = deflateInit(&zstream_, Z_DEFAULT_COMPRESSION);
}

~ZlibOutputStream()
{
finish();
}

// 返回错误信息
const char* zlibErrorMessage() const { return zstream_.msg; }

// 返回错误码
int zlibErrorCode() const { return zerror_; }

// 返回输入流的字节数
int64_t inputBytes() const { return zstream_.total_in; }

// 返回输出流的字节数
int64_t outputBytes() const { return zstream_.total_out; }

// 内部缓冲区的大小
int internalOutputBufferSize() const { return bufferSize_; }

// 写入数据，然后进行压缩
bool write(StringPiece buf)
{
if (zerror_ != Z_OK)
return false;

assert(zstream_.next_in == NULL && zstream_.avail_in == 0);

void* in = const_cast<char*>(buf.data());

// 设置压缩数据起始位置和长度
zstream_.next_in = static_cast<Bytef*>(in);
zstream_.avail_in = buf.size();

// 压缩
while (zstream_.avail_in > 0 && zerror_ == Z_OK)
{
zerror_ = compress(Z_NO_FLUSH);
}

if (zstream_.avail_in == 0)
{
assert(static_cast<const void*>(zstream_.next_in) == buf.end());
zstream_.next_in = NULL;
}
return zerror_ == Z_OK;
}

// 输入数据并进行压缩
bool write(Buffer* input)
{
if (zerror_ != Z_OK)
return false;

void* in = const_cast<char*>(input->peek());

zstream_.next_in = static_cast<Bytef*>(in);
zstream_.avail_in = static_cast<int>(input->readableBytes());

// 压缩
if (zstream_.avail_in > 0 && zerror_ == Z_OK)
{
zerror_ = compress(Z_NO_FLUSH);
}
input->retrieve(input->readableBytes() - zstream_.avail_in);
return zerror_ == Z_OK;
}

// 刷新
bool finish()
{
if (zerror_ != Z_OK)
return false;

while (zerror_ == Z_OK)
{
// 冲刷
zerror_ = compress(Z_FINISH);
}
zerror_ = deflateEnd(&zstream_);
bool ok = zerror_ == Z_OK;
zerror_ = Z_STREAM_END;
return ok;
}

private:

// 压缩
int compress(int flush)
{
output_->ensureWritableBytes(bufferSize_);

zstream_.next_out = reinterpret_cast<Bytef*>(output_->beginWrite());
zstream_.avail_out = static_cast<int>(output_->writableBytes());

// 压缩数据
int error = ::deflate(&zstream_, flush);

output_->hasWritten(output_->writableBytes() - zstream_.avail_out);
if (output_->writableBytes() == 0 && bufferSize_ < 65536)
{
bufferSize_ *= 2;
}
return error;
}

// 输出缓冲区
Buffer* output_;

// 压缩的数据流
z_stream zstream_;

// 错误码
int zerror_;

// 缓冲区的大小
int bufferSize_;
};