zstream: track and limit memory use

`zstream` reads stream record payloads into memory while processing a
stream. Most subcommands use several parallel queues, which currently
have queue lengths on the order of 1000 slots. Depending on the
pipeline, there may be several thousand records in memory at any given
moment.

This is fine for typical filesystem send streams because most records
aren't payload-bearing and default record sizes are modest. However,
it's possible for streams to be far denser, as illustrated by this zvol
send stream breakdown posted by @ryan-moeller in #18565:

```
SUMMARY:
        Total DRR_BEGIN records = 1 (0 bytes)
        Total DRR_END records = 1 (0 bytes)
        Total DRR_OBJECT records = 2 (0 bytes)
        Total DRR_FREEOBJECTS records = 1 (0 bytes)
        Total DRR_WRITE records = 4304609 (70131052544 bytes)
        Total DRR_WRITE_BYREF records = 0 (0 bytes)
        Total DRR_WRITE_EMBEDDED records = 645 (55664 bytes)
        Total DRR_FREE records = 2202 (0 bytes)
        Total DRR_SPILL records = 0 (0 bytes)
        Total DRR_OBJECT_RANGE records = 0 (0 bytes)
        Total DRR_REDACT records = 0 (0 bytes)
        Total records = 4307461
        Total payload size = 70131108208 (0x105423c970)
        Total header overhead = 1343927832 (0x501ab618)
        Total stream length = 71475036040 (0x10a43e7f88)
```

This PR standardizes functions for manipulating payloads within
`zstream` and adds memory accounting and read throttling. Modules
anywhere within a processing chain may want to modify or dispose of
payloads, so memory accounting is thread-safe.

Throttling is applied at the point of ingress, `chain_read()`. That
function blocks when the limit is reached. However, parallel queues and
later chain segments continue to function normally, so data naturally
drains out of the pipeline over time.

The memory budget is currently set at 512MB plus 10% of system memory
beyond 4GB. There's no reason these parameters can't be made
configurable in the future, but for now they're hard-coded. The budget
is just an accounting limit and there is no up-front allocation.
Typical processing of filesystem streams will not approach this limit.

Memory accounting covers only payload data attached to DRR packets
within a processing chain. When data is removed from a stream for
separate handling (now achieved by calling `export_payload()`), it's no
longer tracked by the standard memory accounting.

Reviewed-by: Brian Behlendorf <[email protected]>
Signed-off-by: Garth Snyder <[email protected]>
Closes #18982

cmd/zstream/zstream_decompress.c

@@ -102,9 +102,7 @@ chain_decompress_named_writes(void *item_in, void *context)
 		    (u_longlong_t)drrw->drr_object,
 		    (u_longlong_t)drrw->drr_offset);
 	} else {
-		free(item->dp_payload);
-		item->dp_payload = dcbuff;
-		item->dp_payload_size = drrw->drr_logical_size;
+		set_payload(item, dcbuff, drrw->drr_logical_size);
 		drrw->drr_compressiontype = 0;
 		drrw->drr_compressed_size = 0;
 		if (OPTION_ENABLED(CA_VERBOSE)) {

cmd/zstream/zstream_drop_record.c

@@ -66,22 +66,7 @@ chain_drop_records(void *item_in, void *context)
 			warnx("dropping %s record for object %llu "
 			    "offset %llu", record_type, object, offset);
 		}
-		/*
-		 * It really feels like the chain executor ought to be
-		 * responsible for freeing this payload. However, it
-		 * operates at a more abstract level and knows nothing about
-		 * DMU records and their payloads, so this'll have to be
-		 * done here when the drop decision is made.
-		 *
-		 * Fine for now, but if another case like this comes up in
-		 * the future, the issue probably needs to be handled
-		 * through a more clearly defined path.
-		 */
-		if (item->dp_payload_size && item->dp_payload != NULL) {
-			free(item->dp_payload);
-			item->dp_payload = NULL;
-			item->dp_payload_size = 0;
-		}
+		set_payload(item, NULL, 0);
 		return (D_DROP);
 	}
 

cmd/zstream/zstream_io.c

@@ -17,6 +17,7 @@
 #include <arpa/inet.h>
 #include <err.h>
 #include <libzutil.h>
+#include <pthread.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -28,9 +29,23 @@
 #include <time.h>
 #include <unistd.h>
 
+#include "zstream_chain.h"
 #include "zstream_modules.h"
 #include "zstream_util.h"
 
+/*
+ * Memory devoted to storing payloads is limited to
+ *
+ *   MEMORY_BASE + (system_memory - MEMORY_BASE_CUTOFF) * MEMORY_PCT / 100
+ *
+ * When memory is exhausted, chain_read() waits until memory in use is
+ * MEMORY_HYSTERESIS bytes lower than the nominal limit.
+ */
+#define	MEMORY_BASE		(512 << 20)	/* 512MB */
+#define	MEMORY_BASE_CUTOFF	(4ULL << 30)	/* 4GB */
+#define	MEMORY_PCT		10		/* % beyond the base region */
+#define	MEMORY_HYSTERESIS 	(128 << 20)	/* 128MB */
+
 /*
  * Init only the filename; chain functions will prepare the FILE *
  */
@@ -48,6 +63,21 @@ typedef struct {
 	uint64_t	cc_last_bytes;
 } checkpoint_context_t;
 
+/*
+ * See comments at set_payload() for more information about locking and
+ * performance considerations for data-in-flight tracking. Briefly, access
+ * patterns make the mutex expensive, so it's reserved for awakening threads
+ * that are waiting for memory.
+ */
+typedef struct {
+	pthread_mutex_t	dif_mutex;
+	pthread_cond_t	dif_cond;
+	uint64_t	dif_current;	/* atomic access only */
+	boolean_t	dif_waiting;	/* atomic access only */
+	uint64_t	dif_allowed;
+	uint64_t	dif_resume;	/* dif_allowed - MEMORY_HYSTERESIS */
+} data_in_flight_t;
+
 static io_context_t io_contexts[MAX_IO_STREAMS];
 static int next_io_context = 0;
 
@@ -57,6 +87,33 @@ static int next_checkpoint_context = 0;
 static uint32_t drop_contexts[MAX_DROP_FILTERS];
 static int next_drop_context = 0;
 
+static data_in_flight_t payloads = {
+	.dif_mutex = PTHREAD_MUTEX_INITIALIZER,
+	.dif_cond = PTHREAD_COND_INITIALIZER
+};
+
+static pthread_once_t	dif_init_control = PTHREAD_ONCE_INIT;
+
+/*
+ * Called through setup_io() -> pthread_once()
+ */
+static void
+initialize_memory_tracking(void)
+{
+	ssize_t pagesize = (ssize_t)sysconf(_SC_PAGESIZE);
+	ssize_t pages = (ssize_t)sysconf(_SC_PHYS_PAGES);
+	if (pagesize < 0 || pages < 0) {
+		warnx("unable to read system memory info");
+		payloads.dif_allowed = UINT64_MAX; /* no limit */
+	} else {
+		uint64_t total_mem = (uint64_t)pagesize * (uint64_t)pages;
+		int64_t flex = total_mem - MEMORY_BASE_CUTOFF;
+		int64_t addl = (double)flex * MEMORY_PCT / 100;
+		payloads.dif_allowed = MEMORY_BASE + MAX(addl, 0);
+	}
+	payloads.dif_resume = payloads.dif_allowed - MEMORY_HYSTERESIS;
+}
+
 /*
  * Run from within chain execution to initialize I/O. A NULL filename
  * indicates stdin or stdout.
@@ -210,6 +267,81 @@ set_stream_attributes(drr_packet_t *item)
 	}
 }
 
+/*
+ * Given a desired payload size, determine whether we can read it in
+ * immediately. If not, we have to wait for memory to become available.
+ *
+ * chain_read() is a serial chain step and will always be called by the same
+ * thread. However, multiple other steps in the chain may want to modify or
+ * free payloads, so memory tracking has to be managed with multithreading
+ * in mind.
+ *
+ * The common case is that we are nowhere near the limit, which costs only a
+ * single unsynchronized read of dif_current. We lock the mutex only when we
+ * are actually going to await the condition.
+ *
+ * The store to dif_waiting and the subsequent load of dif_current are the
+ * mirror image of the sequence in set_payload_impl(), which stores
+ * dif_current and then loads dif_waiting. Both halves must be sequentially
+ * consistent: if either were weaker, the two threads could miss each
+ * other's store.
+ */
+static inline void
+maybe_wait_for_memory(size_t bytes_wanted)
+{
+	uint64_t current = __atomic_load_n(&payloads.dif_current,
+	    __ATOMIC_RELAXED);
+
+	if (current + bytes_wanted <= payloads.dif_allowed)
+		return;
+
+	pthread_mutex_lock(&payloads.dif_mutex);
+	__atomic_store_n(&payloads.dif_waiting, B_TRUE, __ATOMIC_SEQ_CST);
+	while (__atomic_load_n(&payloads.dif_current, __ATOMIC_SEQ_CST) >
+	    payloads.dif_resume) {
+		pthread_cond_wait(&payloads.dif_cond, &payloads.dif_mutex);
+	}
+	/*
+	 * A freeing thread that sees a stale B_TRUE here just takes the
+	 * mutex for a broadcast that no one is waiting for, so this store
+	 * needs no ordering of its own.
+	 */
+	__atomic_store_n(&payloads.dif_waiting, B_FALSE, __ATOMIC_RELAXED);
+	pthread_mutex_unlock(&payloads.dif_mutex);
+}
+
+/*
+ * Read in an item's payload. We don't do memory accounting here because
+ * that's now handled by set_payload(). This function reads the payload into
+ * a newly allocated buffer and returns the buffer. set_payload() attaches
+ * an existing buffer to a dp_drr_t item.
+ */
+static inline uint8_t *
+read_payload(io_context_t *context, size_t size)
+{
+	maybe_wait_for_memory(size);
+	uint8_t *buff = safe_malloc(size);
+	size_t n_read = fread(buff, size, 1, context->ic_fp);
+	if (n_read != 1) {
+		if (ferror(context->ic_fp)) {
+			err(1, "error reading record payload at offset %llu",
+			    (u_longlong_t)context->ic_offset);
+		} else {
+			/*
+			 * We can't exit here because ZTS depends on being
+			 * able to process randomly truncated streams.
+			 */
+			warnx("input ends mid-record at offset %llu - "
+			    "stream is likely corrupt",
+			    (u_longlong_t)context->ic_offset);
+			fclose(context->ic_fp);
+			free(buff);
+			return (NULL);
+		}
+	}
+	return (buff);
+}
+
 static disposition_t
 chain_read(void *item_in, void *context_in)
 {
@@ -224,6 +356,10 @@ chain_read(void *item_in, void *context_in)
 	if (!context->ic_fp)
 		open_file(context);
 
+	item->dp_payload = NULL;
+	item->dp_payload_size = 0;
+	item->dp_stream_offset = context->ic_offset;
+
 	if (fread(drr, sizeof (dmu_replay_record_t), 1, context->ic_fp) != 1) {
 		if (ferror(context->ic_fp)) {
 			err(1, "error reading record header at offset %llu",
@@ -236,43 +372,19 @@ chain_read(void *item_in, void *context_in)
 	if (context->ic_offset == 0)
 		set_stream_attributes(item);
 
-	size_t payload_size = calc_payload_size(&item->dp_drr);
+	size_t payload_size = calc_payload_size(drr);
 	if (payload_size > UINT32_MAX) {
-		errx(1, "stated packet size is greater than uint32_t"
+		errx(1, "stated packet size is greater than uint32_t "
 		    "at offset %llu", (u_longlong_t)context->ic_offset);
+	} else if (payload_size > 0) {
+		uint8_t *buff = read_payload(context, payload_size);
+		if (buff == NULL)
+			return (D_EOF);
+		set_payload(item, buff, payload_size);
 	}
-	item->dp_payload_size = payload_size;
-	if (item->dp_payload_size > 0) {
-		item->dp_payload = safe_malloc(item->dp_payload_size);
-		size_t n_read = fread(item->dp_payload, item->dp_payload_size,
-		    1, context->ic_fp);
-		if (n_read != 1) {
-			if (ferror(context->ic_fp)) {
-				err(1, "error reading record payload at "
-				    " offset %llu",
-				    (u_longlong_t)context->ic_offset);
-			} else {
-				/*
-				 * We can't exit here because the ZFS test
-				 * suite depends on being able to process
-				 * streams truncated at random places.
-				 */
-				warnx("input ends mid-record at offset %llu "
-				    "- stream is likely corrupt",
-				    (u_longlong_t)context->ic_offset);
-				fclose(context->ic_fp);
-				free(item->dp_payload);
-				return (D_EOF);
-			}
-		}
-	} else {
-		item->dp_payload = NULL;
-	}
-	item->dp_stream_offset = context->ic_offset;
 
 	uint32_t drr_type = ATTR_IS_SET(CA_BYTESWAPPED) ?
 	    BSWAP_32(drr->drr_type) : drr->drr_type;
-
 	if (drr_type >= DRR_NUMTYPES) {
 		err(1, "invalid record type %llu found at offset %llu",
 		    (u_longlong_t)drr_type, (u_longlong_t)context->ic_offset);
@@ -305,6 +417,8 @@ chain_write(void *item_in, void *context_in)
 				err(1, "error closing output stream");
 			context->ic_fp = NULL;
 		}
+		VERIFY0(__atomic_load_n(&payloads.dif_current,
+		    __ATOMIC_SEQ_CST));
 		return (D_OK);
 	}
 
@@ -321,9 +435,6 @@ chain_write(void *item_in, void *context_in)
 		    item->dp_payload_size, 1, context->ic_fp);
 		if (n_written != 1) {
 			err(1, "error writing payload");
-		} else {
-			free(item->dp_payload);
-			item->dp_payload = NULL;
 		}
 	}
 
@@ -340,6 +451,7 @@ chain_write(void *item_in, void *context_in)
 	stats->rs_total_header_bytes += sizeof (dmu_replay_record_t);
 	stats->rs_total_payload_bytes += item->dp_payload_size;
 
+	set_payload(item, NULL, 0);
 	return (D_OK);
 }
 
@@ -352,11 +464,10 @@ chain_null_output(void *item_in, void *context)
 	(void) context;
 	drr_packet_t *item = (drr_packet_t *)item_in;
 
-	if (item && item->dp_payload != NULL && item->dp_payload_size > 0) {
-		free(item->dp_payload);
-		item->dp_payload = NULL;
-		item->dp_payload_size = 0;
-	}
+	if (item == NULL)
+		return (D_OK);
+
+	set_payload(item, NULL, 0);
 	return (D_OK);
 }
 
@@ -366,6 +477,7 @@ chain_null_output(void *item_in, void *context)
 static chain_step_t
 setup_io(const char *filename, boolean_t for_reading)
 {
+	pthread_once(&dif_init_control, initialize_memory_tracking);
 	int context_num = next_io_context++ % MAX_IO_STREAMS;
 
 	io_context_t context = {
@@ -502,11 +614,7 @@ chain_drop_record_types(void *item_in, void *context_in)
 	}
 
 	if (((UINT32_C(1) << type) & *context) != 0) {
-		if (item->dp_payload != NULL) {
-			free(item->dp_payload);
-			item->dp_payload = NULL;
-			item->dp_payload_size = 0;
-		}
+		set_payload(item, NULL, 0);
 		return (D_DROP);
 	}
 	return (D_OK);
@@ -531,3 +639,70 @@ serial_drop_record_types(uint32_t drop_mask)
 	};
 	return (step);
 }
+
+/*
+ * Every record that carries a payload passes through here at least twice,
+ * once on the way in and once on the way out. Those calls are also made by
+ * different threads, which unfortunately is something of an adversarial
+ * pattern for a pthreads mutex. The cache line containing the lock
+ * structure ping-pongs among cores, and each move is performed with
+ * restrictive memory barriers. We use atomic operations instead and reserve
+ * the mutex for waking up a sleeping memory consumer.
+ *
+ * The atomic update is done in sequentially consistent mode because the
+ * (potential) load of dif_waiting beneath must not be hoisted above the
+ * update to dif_current. dif_waiting shares a cache line with dif_current,
+ * which we have just acquired exclusively, so reading it is essentially
+ * free.
+ */
+static void
+set_payload_impl(void *item_in, void *payload_in, uint64_t size,
+    boolean_t free_old)
+{
+	drr_packet_t *item = (drr_packet_t *)item_in;
+	uint8_t *payload = (uint8_t *)payload_in;
+	VERIFY(payload != NULL || size == 0);
+	VERIFY(payload == NULL || payload != item->dp_payload);
+
+	if (free_old && item->dp_payload != NULL) {
+		free(item->dp_payload);
+	}
+	int64_t delta = (int64_t)size - (int64_t)item->dp_payload_size;
+	VERIFY3U(size, <=, UINT32_MAX);
+	item->dp_payload = payload;
+	item->dp_payload_size = (uint32_t)size;
+
+	/*
+	 * Atomics are nominally unsigned operations, but because of twos
+	 * complement arithmetic it's fine to add a "negative" value.
+	 */
+	uint64_t current = __atomic_add_fetch(&payloads.dif_current, delta,
+	    __ATOMIC_SEQ_CST);
+
+	/*
+	 * Only a net reduction can release a parked reader, and only if it
+	 * brings us back under the hysteresis threshold.
+	 */
+	if (delta < 0 && current <= payloads.dif_resume &&
+	    __atomic_load_n(&payloads.dif_waiting, __ATOMIC_SEQ_CST)) {
+		pthread_mutex_lock(&payloads.dif_mutex);
+		pthread_cond_broadcast(&payloads.dif_cond);
+		pthread_mutex_unlock(&payloads.dif_mutex);
+	}
+}
+
+void
+set_payload(void *item_in, void *payload_in, uint64_t size)
+{
+	set_payload_impl(item_in, payload_in, size, B_TRUE);
+}
+
+/*
+ * Remove a payload from the chain's management without freeing. It's up to
+ * the recipient to free the buffer.
+ */
+void
+export_payload(void *item_in)
+{
+	set_payload_impl(item_in, NULL, 0, B_FALSE);
+}

cmd/zstream/zstream_io.h

@@ -71,6 +71,24 @@ serial_read_stream(const char *filename);
 chain_step_t
 serial_write_stream(const char *filename);
 
+/*
+ * When payloads change (e.g., after being decompressed), this function
+ * should always be used to intermediate. It frees the old payload and
+ * updates the accounting for total data in flight. To free the old payload
+ * without replacing it, just pass in NULL.
+ */
+void
+set_payload(void *item_in, void *payload, uint64_t size);
+
+/*
+ * Sometimes, e.g., in the implementation of "zstream raw", we want to take
+ * a payload buffer out of the chain system and hand its control over to some
+ * other system. We need to update the memory accounting but not attempt to
+ * free the buffer.
+ */
+void
+export_payload(void *item_in);
+
 /*
  * Report throughput periodically
  */

cmd/zstream/zstream_raw.c

@@ -342,7 +342,7 @@ chain_replay_raw(void *item_in, void *context_in)
 			 * The buffer is no longer owned by the chain.  We will
 			 * free it when safe.
 			 */
-			item->dp_payload = NULL;
+			export_payload(item);
 		} else if (context->volume.isreg && context->stream.inprop) {
 			ASSERT0(drrw->drr_offset);
 			context->volume.size = apply_properties(context,
@@ -363,20 +363,15 @@ chain_replay_raw(void *item_in, void *context_in)
 		    &drr->drr_u.drr_write_embedded;
 
 		if (!ctype_is_uncompressed(drrwe->drr_compression)) {
-			uint8_t *buffer = item->dp_payload;
-			uint32_t lsize = drrwe->drr_lsize;
-
-			ASSERT3U(item->dp_payload_size, <=, lsize);
-
-			item->dp_payload = decompress_buffer(buffer,
-			    item->dp_payload_size, lsize,
+			VERIFY3U(item->dp_payload_size, <=, drrwe->drr_lsize);
+			uint8_t *debuff = decompress_buffer(item->dp_payload,
+			    item->dp_payload_size, drrwe->drr_lsize,
 			    drrwe->drr_compression);
-			if (item->dp_payload == NULL)
+			if (debuff == NULL)
 				errx(EXIT_FAILURE,
 				    "decompression failed at offset %llu",
 				    (u_longlong_t)drrwe->drr_offset);
-			item->dp_payload_size = lsize;
-			free(buffer);
+			set_payload(item, debuff, drrwe->drr_lsize);
 		}
 		if (context->stream.inzvol) {
 			buffer_write(context, item->dp_payload,
@@ -385,7 +380,7 @@ chain_replay_raw(void *item_in, void *context_in)
 			 * The buffer is no longer owned by the chain.  We will
 			 * free it when safe.
 			 */
-			item->dp_payload = NULL;
+			export_payload(item);
 		} else if (context->volume.isreg && context->stream.inprop) {
 			ASSERT0(drrwe->drr_offset);
 			context->volume.size = apply_properties(context,

cmd/zstream/zstream_recompress.c

@@ -141,9 +141,7 @@ chain_decompress_writes(queue_item_t *item_in, void *context)
 		    (u_longlong_t)drrw->drr_object,
 		    (u_longlong_t)drrw->drr_offset);
 	}
-	free(item->dp_payload);
-	item->dp_payload = debuff;
-	item->dp_payload_size = drrw->drr_logical_size;
+	set_payload(item, debuff, drrw->drr_logical_size);
 	drrw->drr_compressed_size = 0;
 	drrw->drr_compressiontype = 0;
 }
@@ -174,9 +172,7 @@ chain_compress_writes(queue_item_t *item_in, void *context_in)
 		drrw->drr_compressiontype = 0;
 		drrw->drr_compressed_size = 0;
 	} else {
-		free(item->dp_payload);
-		item->dp_payload = cbuff;
-		item->dp_payload_size = csize;
+		set_payload(item, cbuff, csize);
 		drrw->drr_compressed_size = csize;
 		drrw->drr_compressiontype = context->cs_type;
 	}

cmd/zstream/zstream_redup.c

@@ -153,13 +153,12 @@ chain_redup_writes(void *item_in, void *context_in)
 		VERIFY3U(drrw->drr_object, ==, drrwb.drr_refobject);
 		VERIFY3U(drrw->drr_offset, ==, drrwb.drr_refoffset);
 
-		item->dp_payload_size = DRR_WRITE_PAYLOAD_SIZE(drrw);
-		item->dp_payload = safe_malloc(item->dp_payload_size);
-
-		size_t n_read = fread(item->dp_payload, item->dp_payload_size,
-		    1, context->rc_fp);
+		uint64_t size = DRR_WRITE_PAYLOAD_SIZE(drrw);
+		uint8_t *buff = safe_malloc(size);
+		size_t n_read = fread(buff, size, 1, context->rc_fp);
 		if (n_read != 1)
 			err(1, "read of prior payload failed");
+		set_payload(item, buff, size);
 
 		drrw->drr_toguid = drrwb.drr_toguid;
 		drrw->drr_object = drrwb.drr_object;