fz_store: Reap passes.

A few commits back, we introduced the fz_key_storable concept
to allow us to cope with objects that were used both as values
within the store and as parts of keys within the store.

This commit worked, but showed up performance problems; when the
store has several million PDF objects in it, bulk changes (such
as dropping a display list or document) could trigger many passes
across the store.

We therefore introduce a mechanism to ameliorate this. These
passes, now known as "reap passes", can be batched together using
fz_defer_reap_start and fz_defer_reap_end.

We trigger this start/end around display list dropping, and around
PDF content stream processing. This should be fine, as deferral
will be interrupted if we ever run our of memory during mallocing.

8 files changed, 278 insertions, 66 deletions

diff --git a/include/mupdf/fitz/context.h b/include/mupdf/fitz/context.h
index 50e8eeaa..47aa8fa5 100644
--- a/include/mupdf/fitz/context.h
+++ b/include/mupdf/fitz/context.h
@@ -361,8 +361,8 @@ struct fz_locks_context_s
 };
 
 enum {
-	FZ_LOCK_ALLOC = 0,
-	FZ_LOCK_FILE, /* Unused now */
+	FZ_LOCK_REAP = 0,
+	FZ_LOCK_ALLOC,
 	FZ_LOCK_FREETYPE,
 	FZ_LOCK_GLYPHCACHE,
 	FZ_LOCK_MAX
diff --git a/include/mupdf/fitz/store.h b/include/mupdf/fitz/store.h
index 2ef9810b..a3c45b7f 100644
--- a/include/mupdf/fitz/store.h
+++ b/include/mupdf/fitz/store.h
@@ -37,7 +37,8 @@ struct fz_storable_s {
 
 struct fz_key_storable_s {
 	fz_storable storable;
-	int store_key_refs;
+	short store_key_refs;
+	unsigned short needs_reaping;
 };
 
 #define FZ_INIT_STORABLE(S_,RC,DROP) \
@@ -47,14 +48,14 @@ struct fz_key_storable_s {
 
 #define FZ_INIT_KEY_STORABLE(KS_,RC,DROP) \
 	do { fz_key_storable *KS = &(KS_)->key_storable; KS->store_key_refs = 0;\
-	FZ_INIT_STORABLE(KS,RC,DROP); \
+	KS->needs_reaping = 0; FZ_INIT_STORABLE(KS,RC,DROP); \
 	} while (0)
 
 void *fz_keep_storable(fz_context *, const fz_storable *);
 void fz_drop_storable(fz_context *, const fz_storable *);
 
 void *fz_keep_key_storable(fz_context *, const fz_key_storable *);
-int fz_drop_key_storable(fz_context *, const fz_key_storable *);
+void fz_drop_key_storable(fz_context *, const fz_key_storable *);
 
 void *fz_keep_key_storable_key(fz_context *, const fz_key_storable *);
 void fz_drop_key_storable_key(fz_context *, const fz_key_storable *);
@@ -74,6 +75,39 @@ void fz_drop_key_storable_key(fz_context *, const fz_key_storable *);
 	to an fz_store_hash structure. If make_hash_key function returns 0,
 	then the key is determined not to be hashable, and the value is
 	not stored in the hash table.
+
+	Some objects can be used both as values within the store, and as a
+	component of keys within the store. We refer to these objects as
+	"key storable" objects. In this case, we need to take additional
+	care to ensure that we do not end up keeping an item within the
+	store, purely because it's value is referred to by another key in
+	the store.
+
+	An example of this are fz_images in PDF files. Each fz_image is
+	placed into the	store to enable it to be easily reused. When the
+	image is rendered, a pixmap is generated from the image, and the
+	pixmap is placed into the store so it can be reused on subsequent
+	renders. The image forms part of the key for the pixmap.
+
+	When we close the pdf document (and any associated pages/display
+	lists etc), we drop the images from the store. This may leave us
+	in the position of the images having non-zero reference counts
+	purely because they are used as part of the keys for the pixmaps.
+
+	We therefore use special reference counting functions to keep
+	track of these "key storable" items, and hence store the number of
+	references to these items that are used in keys.
+
+	When the number of references to an object == the number of
+	references to an object from keys in the store, we know that we can
+	remove all the items which have that object as part of the key.
+	This is done by running a pass over the store, 'reaping' those
+	items.
+
+	Reap passes are slower than we would like as they touching every
+	item in the store. We therefore provide a way to 'batch' such
+	reap passes together, using fz_defer_reap_start/fz_defer_reap_end
+	to bracket a region in which many may be triggered.
 */
 typedef struct fz_store_hash_s fz_store_hash;
 
@@ -110,6 +144,7 @@ struct fz_store_type_s
 	void (*drop_key)(fz_context *,void *);
 	int (*cmp_key)(fz_context *ctx, void *, void *);
 	void (*print)(fz_context *ctx, fz_output *out, void *);
+	int (*needs_reap)(fz_context *ctx, void *);
 };
 
 /*
@@ -219,4 +254,33 @@ void fz_filter_store(fz_context *ctx, fz_store_filter_fn *fn, void *arg, fz_stor
 void fz_print_store(fz_context *ctx, fz_output *out);
 void fz_print_store_locked(fz_context *ctx, fz_output *out);
 
+/*
+	fz_defer_reap_start: Increment the defer reap count.
+
+	No reap operations will take place (except for those
+	triggered by an immediate failed malloc) until the
+	defer reap count returns to 0.
+
+	Call this at the start of a process during which you
+	potentially might drop many reapable objects.
+
+	It is vital that every fz_defer_reap_start is matched
+	by a fz_defer_reap_end call.
+*/
+void fz_defer_reap_start(fz_context *ctx);
+
+/*
+	fz_defer_reap_end: Decrement the defer reap count.
+
+	If the defer reap count returns to 0, and the store
+	has reapable objects in, a reap pass will begin.
+
+	Call this at the end of a process during which you
+	potentially might drop many reapable objects.
+
+	It is vital that every fz_defer_reap_start is matched
+	by a fz_defer_reap_end call.
+*/
+void fz_defer_reap_end(fz_context *ctx);
+
 #endif
diff --git a/source/fitz/image.c b/source/fitz/image.c
index 247cb088..8d692fbd 100644
--- a/source/fitz/image.c
+++ b/source/fitz/image.c
@@ -82,39 +82,34 @@ fz_cmp_image_key(fz_context *ctx, void *k0_, void *k1_)
 }
 
 static void
-fz_print_image(fz_context *ctx, fz_output *out, void *key_)
+fz_print_image_key(fz_context *ctx, fz_output *out, void *key_)
 {
 	fz_image_key *key = (fz_image_key *)key_;
 	fz_printf(ctx, out, "(image %d x %d sf=%d) ", key->image->w, key->image->h, key->l2factor);
 }
 
+static int
+fz_needs_reap_image_key(fz_context *ctx, void *key_)
+{
+	fz_image_key *key = (fz_image_key *)key_;
+
+	return (key->image->key_storable.needs_reaping);
+}
+
 static fz_store_type fz_image_store_type =
 {
 	fz_make_hash_image_key,
 	fz_keep_image_key,
 	fz_drop_image_key,
 	fz_cmp_image_key,
-	fz_print_image
+	fz_print_image_key,
+	fz_needs_reap_image_key
 };
 
-static int
-drop_matching_images(fz_context *ctx, void *image_, void *key_)
-{
-	fz_image_key *key = (fz_image_key *)key_;
-	fz_image *image = (fz_image *)image_;
-
-	return key->image == image;
-}
-
 void
 fz_drop_image(fz_context *ctx, fz_image *image)
 {
-	if (fz_drop_key_storable(ctx, &image->key_storable))
-	{
-		/* All the image refs left are references from keys in the store. */
-		/* We can never hope to match these keys again, so drop the objects. */
-		fz_filter_store(ctx, drop_matching_images, image, &fz_image_store_type);
-	}
+	fz_drop_key_storable(ctx, &image->key_storable);
 }
 
 static void
diff --git a/source/fitz/list-device.c b/source/fitz/list-device.c
index dc8b1f0b..3d17ee7d 100644
--- a/source/fitz/list-device.c
+++ b/source/fitz/list-device.c
@@ -1369,7 +1369,9 @@ fz_keep_display_list(fz_context *ctx, fz_display_list *list)
 void
 fz_drop_display_list(fz_context *ctx, fz_display_list *list)
 {
+	fz_defer_reap_start(ctx);
 	fz_drop_storable(ctx, &list->storable);
+	fz_defer_reap_end(ctx);
 }
 
 fz_rect *
diff --git a/source/fitz/store.c b/source/fitz/store.c
index d9674eb6..639e0662 100644
--- a/source/fitz/store.c
+++ b/source/fitz/store.c
@@ -29,6 +29,10 @@ struct fz_store_s
 	/* We keep track of the size of the store, and keep it below max. */
 	size_t max;
 	size_t size;
+
+	/* Protected by the reap lock */
+	int defer_reap_count;
+	int needs_reaping;
 };
 
 void
@@ -50,6 +54,8 @@ fz_new_store_context(fz_context *ctx, size_t max)
 	store->tail = NULL;
 	store->size = 0;
 	store->max = max;
+	store->defer_reap_count = 0;
+	store->needs_reaping = 0;
 	ctx->store = store;
 }
 
@@ -86,16 +92,93 @@ void *fz_keep_key_storable(fz_context *ctx, const fz_key_storable *sc)
 	return fz_keep_storable(ctx, &sc->storable);
 }
 
-int fz_drop_key_storable(fz_context *ctx, const fz_key_storable *sc)
+/*
+	Entered with FZ_LOCK_ALLOC and FZ_LOCK_REAP held.
+	Drops FZ_LOCK_ALLOC.
+*/
+static void
+do_reap(fz_context *ctx)
+{
+	fz_store *store = ctx->store;
+	fz_item *item, *prev, *remove;
+
+	if (store == NULL)
+	{
+		fz_unlock(ctx, FZ_LOCK_ALLOC);
+		return;
+	}
+
+	fz_assert_lock_held(ctx, FZ_LOCK_ALLOC);
+	fz_assert_lock_held(ctx, FZ_LOCK_REAP);
+
+	ctx->store->needs_reaping = 0;
+
+	/* Reap the items */
+	remove = NULL;
+	for (item = store->tail; item; item = prev)
+	{
+		prev = item->prev;
+
+		if (item->type->needs_reap == NULL || item->type->needs_reap(ctx, item->key) == 0)
+			continue;
+
+		/* We have to drop it */
+		store->size -= item->size;
+
+		/* Unlink from the linked list */
+		if (item->next)
+			item->next->prev = item->prev;
+		else
+			store->tail = item->prev;
+		if (item->prev)
+			item->prev->next = item->next;
+		else
+			store->head = item->next;
+
+		/* Remove from the hash table */
+		if (item->type->make_hash_key)
+		{
+			fz_store_hash hash = { NULL };
+			hash.drop = item->val->drop;
+			if (item->type->make_hash_key(ctx, &hash, item->key))
+				fz_hash_remove(ctx, store->hash, &hash);
+		}
+
+		/* Store whether to drop this value or not in 'prev' */
+		item->prev = (item->val->refs > 0 && --item->val->refs == 0) ? item : NULL;
+
+		/* Store it in our removal chain - just singly linked */
+		item->next = remove;
+		remove = item;
+	}
+	fz_unlock(ctx, FZ_LOCK_ALLOC);
+
+	/* Now drop the remove chain */
+	for (item = remove; item != NULL; item = remove)
+	{
+		remove = item->next;
+
+		/* Drop a reference to the value (freeing if required) */
+		if (item->prev)
+			item->val->drop(ctx, item->val);
+
+		/* Always drops the key and drop the item */
+		item->type->drop_key(ctx, item->key);
+		fz_free(ctx, item);
+	}
+
+}
+
+void fz_drop_key_storable(fz_context *ctx, const fz_key_storable *sc)
 {
 	/* Explicitly drop const to allow us to use const
 	 * sanely throughout the code. */
 	fz_key_storable *s = (fz_key_storable *)sc;
 	int drop;
-	int only_refs_are_store_key_refs = 0;
+	int unlock = 1;
 
 	if (s == NULL)
-		return 0;
+		return;
 
 	if (s->storable.refs > 0)
 		(void)Memento_dropRef(s);
@@ -104,11 +187,22 @@ int fz_drop_key_storable(fz_context *ctx, const fz_key_storable *sc)
 	{
 		drop = --s->storable.refs == 0;
 		if (!drop && s->storable.refs == s->store_key_refs)
-			only_refs_are_store_key_refs = 1;
+		{
+			fz_lock(ctx, FZ_LOCK_REAP);
+			if (ctx->store->defer_reap_count > 0)
+				ctx->store->needs_reaping = 1;
+			else
+			{
+				do_reap(ctx);
+				unlock = 0;
+			}
+			fz_unlock(ctx, FZ_LOCK_REAP);
+		}
 	}
 	else
 		drop = 0;
-	fz_unlock(ctx, FZ_LOCK_ALLOC);
+	if (unlock)
+		fz_unlock(ctx, FZ_LOCK_ALLOC);
 	/*
 		If we are dropping the last reference to an object, then
 		it cannot possibly be in the store (as the store always
@@ -118,7 +212,6 @@ int fz_drop_key_storable(fz_context *ctx, const fz_key_storable *sc)
 	 */
 	if (drop)
 		s->storable.drop(ctx, &s->storable);
-	return only_refs_are_store_key_refs;
 }
 
 void *fz_keep_key_storable_key(fz_context *ctx, const fz_key_storable *sc)
@@ -385,8 +478,25 @@ fz_store_item(fz_context *ctx, void *key, void *val_, size_t itemsize, fz_store_
 		size = store->size + itemsize;
 		while (size > store->max)
 		{
+			size_t saved;
+			int relock = 0;
+
+			/* First, do any outstanding reaping, even if defer_reap_count > 0 */
+			fz_lock(ctx, FZ_LOCK_REAP);
+			if (store->needs_reaping)
+			{
+				do_reap(ctx); /* Drops alloc lock */
+				relock = 1;
+			}
+			fz_unlock(ctx, FZ_LOCK_REAP);
+			if (relock)
+				fz_lock(ctx, FZ_LOCK_ALLOC);
+			size = store->size + itemsize;
+			if (size <= store->max)
+				break;
+
 			/* ensure_space may drop, then retake the lock */
-			size_t saved = ensure_space(ctx, size - store->max);
+			saved = ensure_space(ctx, size - store->max);
 			size -= saved;
 			if (saved == 0)
 			{
@@ -794,5 +904,32 @@ void fz_filter_store(fz_context *ctx, fz_store_filter_fn *fn, void *arg, fz_stor
 		item->type->drop_key(ctx, item->key);
 		fz_free(ctx, item);
 	}
+}
+
+void fz_defer_reap_start(fz_context *ctx)
+{
+	if (ctx == NULL || ctx->store == NULL)
+		return;
+
+	fz_lock(ctx, FZ_LOCK_REAP);
+	ctx->store->defer_reap_count++;
+	fz_unlock(ctx, FZ_LOCK_REAP);
+}
+
+void fz_defer_reap_end(fz_context *ctx)
+{
+	int reap;
+
+	if (ctx == NULL || ctx->store == NULL)
+		return;
 
+	fz_lock(ctx, FZ_LOCK_ALLOC);
+	fz_lock(ctx, FZ_LOCK_REAP);
+	--ctx->store->defer_reap_count;
+	reap = ctx->store->defer_reap_count == 0 && ctx->store->needs_reaping;
+	if (reap)
+		do_reap(ctx); /* Drops FZ_LOCK_ALLOC */
+	fz_unlock(ctx, FZ_LOCK_REAP);
+	if (!reap)
+		fz_unlock(ctx, FZ_LOCK_ALLOC);
 }
diff --git a/source/pdf/pdf-interpret.c b/source/pdf/pdf-interpret.c
index f3840794..00c4176c 100644
--- a/source/pdf/pdf-interpret.c
+++ b/source/pdf/pdf-interpret.c
@@ -1231,12 +1231,14 @@ pdf_process_contents(fz_context *ctx, pdf_processor *proc, pdf_document *doc, pd
 
 	fz_try(ctx)
 	{
+		fz_defer_reap_start(ctx);
 		stm = pdf_open_contents_stream(ctx, doc, stmobj);
 		pdf_process_stream(ctx, proc, &csi, stm);
 		pdf_process_end(ctx, proc, &csi);
 	}
 	fz_always(ctx)
 	{
+		fz_defer_reap_end(ctx);
 		fz_drop_stream(ctx, stm);
 		pdf_clear_stack(ctx, &csi);
 		pdf_lexbuf_fin(ctx, &buf);
diff --git a/source/pdf/pdf-store.c b/source/pdf/pdf-store.c
index 44575691..33d34ab9 100644
--- a/source/pdf/pdf-store.c
+++ b/source/pdf/pdf-store.c
@@ -47,7 +47,8 @@ static fz_store_type pdf_obj_store_type =
 	pdf_keep_key,
 	pdf_drop_key,
 	pdf_cmp_key,
-	pdf_print_key
+	pdf_print_key,
+	NULL
 };
 
 void
diff --git a/source/pdf/pdf-xref.c b/source/pdf/pdf-xref.c
index c4db8e4f..7d21775a 100644
--- a/source/pdf/pdf-xref.c
+++ b/source/pdf/pdf-xref.c
@@ -1569,54 +1569,65 @@ pdf_drop_document_imp(fz_context *ctx, pdf_document *doc)
 	if (!doc)
 		return;
 
-	/* Type3 glyphs in the glyph cache can contain pdf_obj pointers
-	 * that we are about to destroy. Simplest solution is to bin the
-	 * glyph cache at this point. */
-	fz_purge_glyph_cache(ctx);
+	fz_try(ctx)
+	{
+		fz_defer_reap_start(ctx);
 
-	if (doc->js)
-		pdf_drop_js(ctx, doc->js);
+		/* Type3 glyphs in the glyph cache can contain pdf_obj pointers
+		 * that we are about to destroy. Simplest solution is to bin the
+		 * glyph cache at this point. */
+		fz_purge_glyph_cache(ctx);
 
-	pdf_drop_xref_sections(ctx, doc);
-	fz_free(ctx, doc->xref_index);
+		if (doc->js)
+			pdf_drop_js(ctx, doc->js);
 
-	if (doc->focus_obj)
-		pdf_drop_obj(ctx, doc->focus_obj);
-	if (doc->file)
-		fz_drop_stream(ctx, doc->file);
-	if (doc->crypt)
-		pdf_drop_crypt(ctx, doc->crypt);
+		pdf_drop_xref_sections(ctx, doc);
+		fz_free(ctx, doc->xref_index);
 
-	pdf_drop_obj(ctx, doc->linear_obj);
-	if (doc->linear_page_refs)
-	{
-		for (i=0; i < doc->page_count; i++)
+		if (doc->focus_obj)
+			pdf_drop_obj(ctx, doc->focus_obj);
+		if (doc->file)
+			fz_drop_stream(ctx, doc->file);
+		if (doc->crypt)
+			pdf_drop_crypt(ctx, doc->crypt);
+
+		pdf_drop_obj(ctx, doc->linear_obj);
+		if (doc->linear_page_refs)
 		{
-			pdf_drop_obj(ctx, doc->linear_page_refs[i]);
+			for (i=0; i < doc->page_count; i++)
+			{
+				pdf_drop_obj(ctx, doc->linear_page_refs[i]);
+			}
+			fz_free(ctx, doc->linear_page_refs);
 		}
-		fz_free(ctx, doc->linear_page_refs);
-	}
-	fz_free(ctx, doc->hint_page);
-	fz_free(ctx, doc->hint_shared_ref);
-	fz_free(ctx, doc->hint_shared);
-	fz_free(ctx, doc->hint_obj_offsets);
+		fz_free(ctx, doc->hint_page);
+		fz_free(ctx, doc->hint_shared_ref);
+		fz_free(ctx, doc->hint_shared);
+		fz_free(ctx, doc->hint_obj_offsets);
 
-	for (i=0; i < doc->num_type3_fonts; i++)
-	{
-		fz_decouple_type3_font(ctx, doc->type3_fonts[i], (void *)doc);
-		fz_drop_font(ctx, doc->type3_fonts[i]);
-	}
-	fz_free(ctx, doc->type3_fonts);
+		for (i=0; i < doc->num_type3_fonts; i++)
+		{
+			fz_decouple_type3_font(ctx, doc->type3_fonts[i], (void *)doc);
+			fz_drop_font(ctx, doc->type3_fonts[i]);
+		}
+		fz_free(ctx, doc->type3_fonts);
 
-	pdf_drop_ocg(ctx, doc->ocg);
+		pdf_drop_ocg(ctx, doc->ocg);
 
-	pdf_empty_store(ctx, doc);
+		pdf_empty_store(ctx, doc);
 
-	pdf_lexbuf_fin(ctx, &doc->lexbuf.base);
+		pdf_lexbuf_fin(ctx, &doc->lexbuf.base);
 
-	pdf_drop_resource_tables(ctx, doc);
+		pdf_drop_resource_tables(ctx, doc);
 
-	fz_free(ctx, doc);
+		fz_free(ctx, doc);
+	}
+	fz_always(ctx)
+	{
+		fz_defer_reap_end(ctx);
+	}
+	fz_catch(ctx)
+		fz_rethrow(ctx);
 }
 
 void