lib-compression: Implement the new ostream.get_buffer_used/avail_size…

…() APIs

The avail_size() implementation isn't fully correct for bzlib/zlib/lzma.
Fixing it requires larger changes though.

src/lib-compression/ostream-bzlib.c

@@ -152,6 +152,28 @@ static int o_stream_bzlib_flush(struct ostream_private *stream)
 	return o_stream_flush_parent(stream);
 }
 
+static size_t
+o_stream_bzlib_get_buffer_used_size(const struct ostream_private *stream)
+{
+	const struct bzlib_ostream *zstream =
+		(const struct bzlib_ostream *)stream;
+
+	/* outbuf has already compressed data that we're trying to send to the
+	   parent stream. We're not including bzlib's internal compression
+	   buffer size. */
+	return (zstream->outbuf_used - zstream->outbuf_offset) +
+		o_stream_get_buffer_used_size(stream->parent);
+}
+
+static size_t
+o_stream_bzlib_get_buffer_avail_size(const struct ostream_private *stream)
+{
+	/* FIXME: not correct - this is counting compressed size, which may be
+	   too larger than uncompressed size in some situations. Fixing would
+	   require some kind of additional buffering. */
+	return o_stream_get_buffer_avail_size(stream->parent);
+}
+
 static ssize_t
 o_stream_bzlib_sendv(struct ostream_private *stream,
 		    const struct const_iovec *iov, unsigned int iov_count)
@@ -193,6 +215,10 @@ struct ostream *o_stream_create_bz2(struct ostream *output, int level)
 	zstream = i_new(struct bzlib_ostream, 1);
 	zstream->ostream.sendv = o_stream_bzlib_sendv;
 	zstream->ostream.flush = o_stream_bzlib_flush;
+	zstream->ostream.get_buffer_used_size =
+		o_stream_bzlib_get_buffer_used_size;
+	zstream->ostream.get_buffer_avail_size =
+		o_stream_bzlib_get_buffer_avail_size;
 	zstream->ostream.iostream.close = o_stream_bzlib_close;
 
 	ret = BZ2_bzCompressInit(&zstream->zs, level, 0, 0);

src/lib-compression/ostream-lz4.c

@@ -146,6 +146,31 @@ static int o_stream_lz4_flush(struct ostream_private *stream)
 	return o_stream_flush_parent(stream);
 }
 
+static size_t
+o_stream_lz4_get_buffer_used_size(const struct ostream_private *stream)
+{
+	const struct lz4_ostream *zstream =
+		(const struct lz4_ostream *)stream;
+
+	/* outbuf has already compressed data that we're trying to send to the
+	   parent stream. compressbuf isn't included in the return value,
+	   because it needs to be filled up or flushed. */
+	return (zstream->outbuf_used - zstream->outbuf_offset) +
+		o_stream_get_buffer_used_size(stream->parent);
+}
+
+static size_t
+o_stream_lz4_get_buffer_avail_size(const struct ostream_private *stream)
+{
+	const struct lz4_ostream *zstream =
+		(const struct lz4_ostream *)stream;
+
+	/* We're only guaranteed to accept data to compressbuf. The parent
+	   stream might have space, but since compressed data gets written
+	   there it's not really known how much we can actually write there. */
+	return sizeof(zstream->compressbuf) - zstream->compressbuf_offset;
+}
+
 static ssize_t
 o_stream_lz4_sendv(struct ostream_private *stream,
 		    const struct const_iovec *iov, unsigned int iov_count)
@@ -183,6 +208,10 @@ struct ostream *o_stream_create_lz4(struct ostream *output, int level)
 	zstream = i_new(struct lz4_ostream, 1);
 	zstream->ostream.sendv = o_stream_lz4_sendv;
 	zstream->ostream.flush = o_stream_lz4_flush;
+	zstream->ostream.get_buffer_used_size =
+		o_stream_lz4_get_buffer_used_size;
+	zstream->ostream.get_buffer_avail_size =
+		o_stream_lz4_get_buffer_avail_size;
 	zstream->ostream.iostream.close = o_stream_lz4_close;
 
 	i_assert(sizeof(zstream->outbuf) >= sizeof(*hdr));

src/lib-compression/ostream-lzma.c

@@ -160,6 +160,28 @@ static int o_stream_lzma_flush(struct ostream_private *stream)
 	return o_stream_flush_parent(stream);
 }
 
+static size_t
+o_stream_lzma_get_buffer_used_size(const struct ostream_private *stream)
+{
+	const struct lzma_ostream *zstream =
+		(const struct lzma_ostream *)stream;
+
+	/* outbuf has already compressed data that we're trying to send to the
+	   parent stream. We're not including lzma's internal compression
+	   buffer size. */
+	return (zstream->outbuf_used - zstream->outbuf_offset) +
+		o_stream_get_buffer_used_size(stream->parent);
+}
+
+static size_t
+o_stream_lzma_get_buffer_avail_size(const struct ostream_private *stream)
+{
+	/* FIXME: not correct - this is counting compressed size, which may be
+	   too larger than uncompressed size in some situations. Fixing would
+	   require some kind of additional buffering. */
+	return o_stream_get_buffer_avail_size(stream->parent);
+}
+
 static ssize_t
 o_stream_lzma_sendv(struct ostream_private *stream,
 		    const struct const_iovec *iov, unsigned int iov_count)
@@ -201,6 +223,10 @@ struct ostream *o_stream_create_lzma(struct ostream *output, int level)
 	zstream = i_new(struct lzma_ostream, 1);
 	zstream->ostream.sendv = o_stream_lzma_sendv;
 	zstream->ostream.flush = o_stream_lzma_flush;
+	zstream->ostream.get_buffer_used_size =
+		o_stream_lzma_get_buffer_used_size;
+	zstream->ostream.get_buffer_avail_size =
+		o_stream_lzma_get_buffer_avail_size;
 	zstream->ostream.iostream.close = o_stream_lzma_close;
 
 	ret = lzma_easy_encoder(&zstream->strm, level, LZMA_CHECK_CRC64);

src/lib-compression/ostream-zlib.c

@@ -237,6 +237,28 @@ static int o_stream_zlib_flush(struct ostream_private *stream)
 	return o_stream_flush_parent(stream);
 }
 
+static size_t
+o_stream_zlib_get_buffer_used_size(const struct ostream_private *stream)
+{
+	const struct zlib_ostream *zstream =
+		(const struct zlib_ostream *)stream;
+
+	/* outbuf has already compressed data that we're trying to send to the
+	   parent stream. We're not including zlib's internal compression
+	   buffer size. */
+	return (zstream->outbuf_used - zstream->outbuf_offset) +
+		o_stream_get_buffer_used_size(stream->parent);
+}
+
+static size_t
+o_stream_zlib_get_buffer_avail_size(const struct ostream_private *stream)
+{
+	/* FIXME: not correct - this is counting compressed size, which may be
+	   too larger than uncompressed size in some situations. Fixing would
+	   require some kind of additional buffering. */
+	return o_stream_get_buffer_avail_size(stream->parent);
+}
+
 static ssize_t
 o_stream_zlib_sendv(struct ostream_private *stream,
 		    const struct const_iovec *iov, unsigned int iov_count)
@@ -299,6 +321,10 @@ o_stream_create_zlib(struct ostream *output, int level, bool gz)
 	zstream = i_new(struct zlib_ostream, 1);
 	zstream->ostream.sendv = o_stream_zlib_sendv;
 	zstream->ostream.flush = o_stream_zlib_flush;
+	zstream->ostream.get_buffer_used_size =
+		o_stream_zlib_get_buffer_used_size;
+	zstream->ostream.get_buffer_avail_size =
+		o_stream_zlib_get_buffer_avail_size;
 	zstream->ostream.iostream.close = o_stream_zlib_close;
 	zstream->crc = 0;
 	zstream->gz = gz;