[v2,05/10] reftable/block: move ownership of block reader into `struct table_iter`

Commit Message

Patrick Steinhardt April 8, 2024, 12:16 p.m. UTC

The table iterator allows the caller to iterate through all records in a
reftable table. To do so it iterates through all blocks of the desired
type one by one, where for each block it creates a new block iterator
and yields all its entries.

One of the things that is somewhat confusing in this context is who owns
the block reader that is being used to read the blocks and pass them to
the block iterator. Intuitively, as the table iterator is responsible
for iterating through the blocks, one would assume that this iterator is
also responsible for managing the lifecycle of the reader. And while it
somewhat is, the block reader is ultimately stored inside of the block
iterator.

Refactor the code such that the block reader is instead fully managed by
the table iterator. Instead of passing the reader to the block iterator,
we now only end up passing the block data to it. Despite clearing up the
lifecycle of the reader, it will also allow for better reuse of the
reader in subsequent patches.

The following benchmark prints a single matching ref out of 1 million
refs. Before:

  HEAP SUMMARY:
      in use at exit: 13,603 bytes in 125 blocks
    total heap usage: 6,607 allocs, 6,482 frees, 509,635 bytes allocated

After:

  HEAP SUMMARY:
      in use at exit: 13,603 bytes in 125 blocks
    total heap usage: 7,235 allocs, 7,110 frees, 301,481 bytes allocated

Note that while there are more allocation and free calls now, the
overall number of bytes allocated is significantly lower. The number of
allocations will be reduced significantly by the next patch though.

Signed-off-by: Patrick Steinhardt <[e-mail]>
---
 reftable/block.c  |  43 ++++++++++------
 reftable/block.h  |  17 ++++---
 reftable/reader.c | 123 ++++++++++++++++++++++------------------------
 3 files changed, 100 insertions(+), 83 deletions(-)

Patch

diff --git a/reftable/block.c b/reftable/block.c
index fe836c21e5..2d8d0668b3 100644
--- a/reftable/block.c
+++ b/reftable/block.c
@@ -261,12 +261,12 @@  void block_reader_release(struct block_reader *br)
 	reftable_block_done(&br->block);
 }
 
-uint8_t block_reader_type(struct block_reader *r)
+uint8_t block_reader_type(const struct block_reader *r)
 {
 	return r->block.data[r->header_off];
 }
 
-int block_reader_first_key(struct block_reader *br, struct strbuf *key)
+int block_reader_first_key(const struct block_reader *br, struct strbuf *key)
 {
 	int off = br->header_off + 4, n;
 	struct string_view in = {
@@ -286,14 +286,16 @@  int block_reader_first_key(struct block_reader *br, struct strbuf *key)
 	return 0;
 }
 
-static uint32_t block_reader_restart_offset(struct block_reader *br, int i)
+static uint32_t block_reader_restart_offset(const struct block_reader *br, int i)
 {
 	return get_be24(br->restart_bytes + 3 * i);
 }
 
-void block_iter_seek_start(struct block_iter *it, struct block_reader *br)
+void block_iter_seek_start(struct block_iter *it, const struct block_reader *br)
 {
-	it->br = br;
+	it->block = br->block.data;
+	it->block_len = br->block_len;
+	it->hash_size = br->hash_size;
 	strbuf_reset(&it->last_key);
 	it->next_off = br->header_off + 4;
 }
@@ -301,7 +303,7 @@  void block_iter_seek_start(struct block_iter *it, struct block_reader *br)
 struct restart_needle_less_args {
 	int error;
 	struct strbuf needle;
-	struct block_reader *reader;
+	const struct block_reader *reader;
 };
 
 static int restart_needle_less(size_t idx, void *_args)
@@ -340,9 +342,11 @@  static int restart_needle_less(size_t idx, void *_args)
 	return args->needle.len < suffix_len;
 }
 
-void block_iter_copy_from(struct block_iter *dest, struct block_iter *src)
+void block_iter_copy_from(struct block_iter *dest, const struct block_iter *src)
 {
-	dest->br = src->br;
+	dest->block = src->block;
+	dest->block_len = src->block_len;
+	dest->hash_size = src->hash_size;
 	dest->next_off = src->next_off;
 	strbuf_reset(&dest->last_key);
 	strbuf_addbuf(&dest->last_key, &src->last_key);
@@ -351,14 +355,14 @@  void block_iter_copy_from(struct block_iter *dest, struct block_iter *src)
 int block_iter_next(struct block_iter *it, struct reftable_record *rec)
 {
 	struct string_view in = {
-		.buf = it->br->block.data + it->next_off,
-		.len = it->br->block_len - it->next_off,
+		.buf = (unsigned char *) it->block + it->next_off,
+		.len = it->block_len - it->next_off,
 	};
 	struct string_view start = in;
 	uint8_t extra = 0;
 	int n = 0;
 
-	if (it->next_off >= it->br->block_len)
+	if (it->next_off >= it->block_len)
 		return 1;
 
 	n = reftable_decode_key(&it->last_key, &extra, in);
@@ -368,7 +372,7 @@  int block_iter_next(struct block_iter *it, struct reftable_record *rec)
 		return REFTABLE_FORMAT_ERROR;
 
 	string_view_consume(&in, n);
-	n = reftable_record_decode(rec, it->last_key, extra, in, it->br->hash_size,
+	n = reftable_record_decode(rec, it->last_key, extra, in, it->hash_size,
 				   &it->scratch);
 	if (n < 0)
 		return -1;
@@ -378,13 +382,22 @@  int block_iter_next(struct block_iter *it, struct reftable_record *rec)
 	return 0;
 }
 
+void block_iter_reset(struct block_iter *it)
+{
+	strbuf_reset(&it->last_key);
+	it->next_off = 0;
+	it->block = NULL;
+	it->block_len = 0;
+	it->hash_size = 0;
+}
+
 void block_iter_close(struct block_iter *it)
 {
 	strbuf_release(&it->last_key);
 	strbuf_release(&it->scratch);
 }
 
-int block_iter_seek_key(struct block_iter *it, struct block_reader *br,
+int block_iter_seek_key(struct block_iter *it, const struct block_reader *br,
 			struct strbuf *want)
 {
 	struct restart_needle_less_args args = {
@@ -436,7 +449,9 @@  int block_iter_seek_key(struct block_iter *it, struct block_reader *br,
 		it->next_off = block_reader_restart_offset(br, i - 1);
 	else
 		it->next_off = br->header_off + 4;
-	it->br = br;
+	it->block = br->block.data;
+	it->block_len = br->block_len;
+	it->hash_size = br->hash_size;
 
 	reftable_record_init(&rec, block_reader_type(br));
 
diff --git a/reftable/block.h b/reftable/block.h
index 601a1e0e89..d733d45ee0 100644
--- a/reftable/block.h
+++ b/reftable/block.h
@@ -84,16 +84,18 @@  int block_reader_init(struct block_reader *br, struct reftable_block *bl,
 void block_reader_release(struct block_reader *br);
 
 /* Returns the block type (eg. 'r' for refs) */
-uint8_t block_reader_type(struct block_reader *r);
+uint8_t block_reader_type(const struct block_reader *r);
 
 /* Decodes the first key in the block */
-int block_reader_first_key(struct block_reader *br, struct strbuf *key);
+int block_reader_first_key(const struct block_reader *br, struct strbuf *key);
 
 /* Iterate over entries in a block */
 struct block_iter {
 	/* offset within the block of the next entry to read. */
 	uint32_t next_off;
-	struct block_reader *br;
+	const unsigned char *block;
+	size_t block_len;
+	int hash_size;
 
 	/* key for last entry we read. */
 	struct strbuf last_key;
@@ -106,17 +108,20 @@  struct block_iter {
 }
 
 /* Position `it` at start of the block */
-void block_iter_seek_start(struct block_iter *it, struct block_reader *br);
+void block_iter_seek_start(struct block_iter *it, const struct block_reader *br);
 
 /* Position `it` to the `want` key in the block */
-int block_iter_seek_key(struct block_iter *it, struct block_reader *br,
+int block_iter_seek_key(struct block_iter *it, const struct block_reader *br,
 			struct strbuf *want);
 
-void block_iter_copy_from(struct block_iter *dest, struct block_iter *src);
+void block_iter_copy_from(struct block_iter *dest, const struct block_iter *src);
 
 /* return < 0 for error, 0 for OK, > 0 for EOF. */
 int block_iter_next(struct block_iter *it, struct reftable_record *rec);
 
+/* Reset the block iterator to pristine state without releasing its memory. */
+void block_iter_reset(struct block_iter *it);
+
 /* deallocate memory for `it`. The block reader and its block is left intact. */
 void block_iter_close(struct block_iter *it);
 
diff --git a/reftable/reader.c b/reftable/reader.c
index f925570bf3..b77b639751 100644
--- a/reftable/reader.c
+++ b/reftable/reader.c
@@ -220,6 +220,7 @@  struct table_iter {
 	struct reftable_reader *r;
 	uint8_t typ;
 	uint64_t block_off;
+	struct block_reader br;
 	struct block_iter bi;
 	int is_finished;
 };
@@ -227,16 +228,6 @@  struct table_iter {
 	.bi = BLOCK_ITER_INIT \
 }
 
-static void table_iter_copy_from(struct table_iter *dest,
-				 struct table_iter *src)
-{
-	dest->r = src->r;
-	dest->typ = src->typ;
-	dest->block_off = src->block_off;
-	dest->is_finished = src->is_finished;
-	block_iter_copy_from(&dest->bi, &src->bi);
-}
-
 static int table_iter_next_in_block(struct table_iter *ti,
 				    struct reftable_record *rec)
 {
@@ -250,14 +241,8 @@  static int table_iter_next_in_block(struct table_iter *ti,
 
 static void table_iter_block_done(struct table_iter *ti)
 {
-	if (!ti->bi.br) {
-		return;
-	}
-	block_reader_release(ti->bi.br);
-	FREE_AND_NULL(ti->bi.br);
-
-	ti->bi.last_key.len = 0;
-	ti->bi.next_off = 0;
+	block_reader_release(&ti->br);
+	block_iter_reset(&ti->bi);
 }
 
 static int32_t extract_block_size(uint8_t *data, uint8_t *typ, uint64_t off,
@@ -321,32 +306,33 @@  int reader_init_block_reader(struct reftable_reader *r, struct block_reader *br,
 	return err;
 }
 
+static void table_iter_close(struct table_iter *ti)
+{
+	table_iter_block_done(ti);
+	block_iter_close(&ti->bi);
+}
+
 static int table_iter_next_block(struct table_iter *dest,
 				 struct table_iter *src)
 {
-	uint64_t next_block_off = src->block_off + src->bi.br->full_block_size;
-	struct block_reader br = { 0 };
-	int err = 0;
+	uint64_t next_block_off = src->block_off + src->br.full_block_size;
+	int err;
 
 	dest->r = src->r;
 	dest->typ = src->typ;
 	dest->block_off = next_block_off;
 
-	err = reader_init_block_reader(src->r, &br, next_block_off, src->typ);
-	if (err > 0) {
+	err = reader_init_block_reader(src->r, &dest->br, next_block_off, src->typ);
+	if (err > 0)
 		dest->is_finished = 1;
-		return 1;
-	}
-	if (err != 0)
+	if (err) {
+		table_iter_block_done(dest);
 		return err;
-	else {
-		struct block_reader *brp =
-			reftable_malloc(sizeof(struct block_reader));
-		*brp = br;
-
-		dest->is_finished = 0;
-		block_iter_seek_start(&dest->bi, brp);
 	}
+
+	dest->is_finished = 0;
+	block_iter_seek_start(&dest->bi, &dest->br);
+
 	return 0;
 }
 
@@ -377,14 +363,13 @@  static int table_iter_next(struct table_iter *ti, struct reftable_record *rec)
 		 * iterator is drained.
 		 */
 		err = table_iter_next_block(&next, ti);
-		table_iter_block_done(ti);
 		if (err) {
 			ti->is_finished = 1;
 			return err;
 		}
 
-		table_iter_copy_from(ti, &next);
-		block_iter_close(&next.bi);
+		table_iter_close(ti);
+		*ti = next;
 	}
 }
 
@@ -393,16 +378,14 @@  static int table_iter_next_void(void *ti, struct reftable_record *rec)
 	return table_iter_next(ti, rec);
 }
 
-static void table_iter_close(void *p)
+static void table_iter_close_void(void *ti)
 {
-	struct table_iter *ti = p;
-	table_iter_block_done(ti);
-	block_iter_close(&ti->bi);
+	table_iter_close(ti);
 }
 
 static struct reftable_iterator_vtable table_iter_vtable = {
 	.next = &table_iter_next_void,
-	.close = &table_iter_close,
+	.close = &table_iter_close_void,
 };
 
 static void iterator_from_table_iter(struct reftable_iterator *it,
@@ -417,19 +400,16 @@  static int reader_table_iter_at(struct reftable_reader *r,
 				struct table_iter *ti, uint64_t off,
 				uint8_t typ)
 {
-	struct block_reader br = { 0 };
-	struct block_reader *brp = NULL;
+	int err;
 
-	int err = reader_init_block_reader(r, &br, off, typ);
+	err = reader_init_block_reader(r, &ti->br, off, typ);
 	if (err != 0)
 		return err;
 
-	brp = reftable_malloc(sizeof(struct block_reader));
-	*brp = br;
 	ti->r = r;
-	ti->typ = block_reader_type(brp);
+	ti->typ = block_reader_type(&ti->br);
 	ti->block_off = off;
-	block_iter_seek_start(&ti->bi, brp);
+	block_iter_seek_start(&ti->bi, &ti->br);
 	return 0;
 }
 
@@ -454,23 +434,34 @@  static int reader_seek_linear(struct table_iter *ti,
 {
 	struct strbuf want_key = STRBUF_INIT;
 	struct strbuf got_key = STRBUF_INIT;
-	struct table_iter next = TABLE_ITER_INIT;
 	struct reftable_record rec;
 	int err = -1;
 
 	reftable_record_init(&rec, reftable_record_type(want));
 	reftable_record_key(want, &want_key);
 
+	/*
+	 * First we need to locate the block that must contain our record. To
+	 * do so we scan through blocks linearly until we find the first block
+	 * whose first key is bigger than our wanted key. Once we have found
+	 * that block we know that the key must be contained in the preceding
+	 * block.
+	 *
+	 * This algorithm is somewhat unfortunate because it means that we
+	 * always have to seek one block too far and then back up. But as we
+	 * can only decode the _first_ key of a block but not its _last_ key we
+	 * have no other way to do this.
+	 */
 	while (1) {
+		struct table_iter next = TABLE_ITER_INIT;
+
 		err = table_iter_next_block(&next, ti);
 		if (err < 0)
 			goto done;
-
-		if (err > 0) {
+		if (err > 0)
 			break;
-		}
 
-		err = block_reader_first_key(next.bi.br, &got_key);
+		err = block_reader_first_key(&next.br, &got_key);
 		if (err < 0)
 			goto done;
 
@@ -480,16 +471,20 @@  static int reader_seek_linear(struct table_iter *ti,
 		}
 
 		table_iter_block_done(ti);
-		table_iter_copy_from(ti, &next);
+		*ti = next;
 	}
 
-	err = block_iter_seek_key(&ti->bi, ti->bi.br, &want_key);
+	/*
+	 * We have located the block that must contain our record, so we seek
+	 * the wanted key inside of it. If the block does not contain our key
+	 * we know that the corresponding record does not exist.
+	 */
+	err = block_iter_seek_key(&ti->bi, &ti->br, &want_key);
 	if (err < 0)
 		goto done;
 	err = 0;
 
 done:
-	block_iter_close(&next.bi);
 	reftable_record_release(&rec);
 	strbuf_release(&want_key);
 	strbuf_release(&got_key);
@@ -508,6 +503,7 @@  static int reader_seek_indexed(struct reftable_reader *r,
 		.u.idx = { .last_key = STRBUF_INIT },
 	};
 	struct table_iter index_iter = TABLE_ITER_INIT;
+	struct table_iter empty = TABLE_ITER_INIT;
 	struct table_iter next = TABLE_ITER_INIT;
 	int err = 0;
 
@@ -549,7 +545,6 @@  static int reader_seek_indexed(struct reftable_reader *r,
 		 * not exist.
 		 */
 		err = table_iter_next(&index_iter, &index_result);
-		table_iter_block_done(&index_iter);
 		if (err != 0)
 			goto done;
 
@@ -558,7 +553,7 @@  static int reader_seek_indexed(struct reftable_reader *r,
 		if (err != 0)
 			goto done;
 
-		err = block_iter_seek_key(&next.bi, next.bi.br, &want_index.u.idx.last_key);
+		err = block_iter_seek_key(&next.bi, &next.br, &want_index.u.idx.last_key);
 		if (err < 0)
 			goto done;
 
@@ -572,18 +567,20 @@  static int reader_seek_indexed(struct reftable_reader *r,
 			break;
 		}
 
-		table_iter_copy_from(&index_iter, &next);
+		table_iter_close(&index_iter);
+		index_iter = next;
+		next = empty;
 	}
 
 	if (err == 0) {
-		struct table_iter empty = TABLE_ITER_INIT;
 		struct table_iter *malloced = reftable_calloc(1, sizeof(*malloced));
-		*malloced = empty;
-		table_iter_copy_from(malloced, &next);
+		*malloced = next;
+		next = empty;
 		iterator_from_table_iter(it, malloced);
 	}
+
 done:
-	block_iter_close(&next.bi);
+	table_iter_close(&next);
 	table_iter_close(&index_iter);
 	reftable_record_release(&want_index);
 	reftable_record_release(&index_result);