--- a/src/plugins/quota/Makefile.am	Sun Dec 03 20:53:45 2006 +0200
+++ b/src/plugins/quota/Makefile.am	Sun Dec 03 20:55:38 2006 +0200
@@ -14,6 +14,7 @@
 
 lib01_quota_plugin_la_SOURCES = \
 	quota.c \
+	quota-count.c \
 	quota-fs.c \
 	quota-dict.c \
 	quota-dirsize.c \

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/src/plugins/quota/quota-count.c	Sun Dec 03 20:55:38 2006 +0200
@@ -0,0 +1,95 @@
+/* Copyright (C) 2006 Timo Sirainen */
+
+#include "lib.h"
+#include "array.h"
+#include "mail-search.h"
+#include "mail-storage.h"
+#include "quota-private.h"
+
+static int quota_count_mailbox(struct mail_storage *storage, const char *name,
+			       uint64_t *bytes_r, uint64_t *count_r)
+{
+	struct mailbox *box;
+	struct mailbox_transaction_context *trans;
+	struct mail_search_context *ctx;
+	struct mail *mail;
+	struct mail_search_arg search_arg;
+	uoff_t size;
+	int ret = 0;
+
+	box = mailbox_open(storage, name, NULL,
+			   MAILBOX_OPEN_READONLY | MAILBOX_OPEN_KEEP_RECENT);
+	if (box == NULL)
+		return -1;
+
+	memset(&search_arg, 0, sizeof(search_arg));
+	search_arg.type = SEARCH_ALL;
+
+	trans = mailbox_transaction_begin(box, 0);
+	ctx = mailbox_search_init(trans, NULL, &search_arg, NULL);
+	mail = mail_alloc(trans, MAIL_FETCH_PHYSICAL_SIZE, NULL);
+	while (mailbox_search_next(ctx, mail) > 0) {
+		size = mail_get_physical_size(mail);
+		if (size != (uoff_t)-1)
+			*bytes_r += size;
+		*count_r += 1;
+	}
+	mail_free(&mail);
+	if (mailbox_search_deinit(&ctx) < 0)
+		ret = -1;
+
+	if (ret < 0)
+		mailbox_transaction_rollback(&trans);
+	else
+		(void)mailbox_transaction_commit(&trans, 0);
+
+	mailbox_close(&box);
+	return ret;
+}
+
+static int quota_count_storage(struct mail_storage *storage,
+			       uint64_t *bytes, uint64_t *count)
+{
+	struct mailbox_list_iterate_context *ctx;
+	struct mailbox_info *info;
+	int ret = 0;
+
+	ctx = mailbox_list_iter_init(storage->list, "*",
+				     MAILBOX_LIST_ITER_FAST_FLAGS);
+	while ((info = mailbox_list_iter_next(ctx)) != NULL) {
+		if ((info->flags & (MAILBOX_NONEXISTENT |
+				    MAILBOX_NOSELECT)) == 0) {
+			ret = quota_count_mailbox(storage, info->name,
+						  bytes, count);
+			if (ret < 0)
+				break;
+		}
+	}
+	if (mailbox_list_iter_deinit(&ctx) < 0)
+		ret = -1;
+
+	return ret;
+}
+
+int quota_count(struct quota *quota, uint64_t *bytes_r, uint64_t *count_r)
+{
+	struct mail_storage *const *storages;
+	unsigned int i, count;
+	int ret;
+
+	i_assert(!quota->counting);
+
+	*bytes_r = *count_r = 0;
+
+	quota->counting = TRUE;
+
+	storages = array_get(&quota->storages, &count);
+	for (i = 0; i < count; i++) {
+		ret = quota_count_storage(storages[i], bytes_r, count_r);
+		if (ret < 0)
+			break;
+	}
+	quota->counting = FALSE;
+
+	return ret;
+}

--- a/src/plugins/quota/quota-dict.c	Sun Dec 03 20:53:45 2006 +0200
+++ b/src/plugins/quota/quota-dict.c	Sun Dec 03 20:55:38 2006 +0200
@@ -72,29 +72,56 @@
 }
 
 static int
+dict_quota_count(struct dict_quota_root *root,
+		 bool want_bytes, uint64_t *value_r)
+{
+	struct dict_transaction_context *dt;
+	uint64_t bytes, count;
+
+	if (quota_count(root->root.quota, &bytes, &count) < 0)
+		return -1;
+
+	t_push();
+	dt = dict_transaction_begin(root->dict);
+	dict_set(dt, DICT_QUOTA_CURRENT_BYTES_PATH, dec2str(bytes));
+	dict_set(dt, DICT_QUOTA_CURRENT_COUNT_PATH, dec2str(count));
+	t_pop();
+
+	if (dict_transaction_commit(dt) < 0)
+		i_error("dict_quota: Couldn't update quota");
+
+	*value_r = want_bytes ? bytes : count;
+	return 1;
+}
+
+static int
 dict_quota_get_resource(struct quota_root *_root, const char *name,
 			uint64_t *value_r, uint64_t *limit __attr_unused__)
 {
 	struct dict_quota_root *root = (struct dict_quota_root *)_root;
 	const char *value;
+	bool want_bytes;
 	int ret;
 
-	if (strcmp(name, QUOTA_NAME_STORAGE_BYTES) == 0) {
-		t_push();
-		ret = dict_lookup(root->dict, unsafe_data_stack_pool,
-				  DICT_QUOTA_CURRENT_BYTES_PATH, &value);
-		*value_r = ret <= 0 ? 0 : strtoull(value, NULL, 10);
-		t_pop();
-	} else if (strcmp(name, QUOTA_NAME_MESSAGES) == 0) {
-		t_push();
-		ret = dict_lookup(root->dict, unsafe_data_stack_pool,
-				  DICT_QUOTA_CURRENT_COUNT_PATH, &value);
-		*value_r = ret <= 0 ? 0 : strtoull(value, NULL, 10);
-		t_pop();
-	} else {
+	if (strcmp(name, QUOTA_NAME_STORAGE_BYTES) == 0)
+		want_bytes = TRUE;
+	else if (strcmp(name, QUOTA_NAME_MESSAGES) == 0)
+		want_bytes = FALSE;
+	else
 		return 0;
-	}
 
+	t_push();
+	ret = dict_lookup(root->dict, unsafe_data_stack_pool,
+			  want_bytes ? DICT_QUOTA_CURRENT_BYTES_PATH :
+			  DICT_QUOTA_CURRENT_COUNT_PATH, &value);
+	if (ret < 0)
+		*value_r = 0;
+	else if (ret == 0)
+		ret = dict_quota_count(root, want_bytes, value_r);
+	else
+		*value_r = strtoull(value, NULL, 10);
+
+	t_pop();
 	return ret;
 }
 
@@ -106,10 +133,14 @@
 	struct dict_transaction_context *dt;
 
 	dt = dict_transaction_begin(root->dict);
-	dict_atomic_inc(dt, DICT_QUOTA_CURRENT_BYTES_PATH,
-			ctx->bytes_used);
-	dict_atomic_inc(dt, DICT_QUOTA_CURRENT_COUNT_PATH,
-			ctx->count_used);
+	if (ctx->bytes_used != 0) {
+		dict_atomic_inc(dt, DICT_QUOTA_CURRENT_BYTES_PATH,
+				ctx->bytes_used);
+	}
+	if (ctx->count_used != 0) {
+		dict_atomic_inc(dt, DICT_QUOTA_CURRENT_COUNT_PATH,
+				ctx->count_used);
+	}
 	
 	if (dict_transaction_commit(dt) < 0)
 		return -1;

--- a/src/plugins/quota/quota-private.h	Sun Dec 03 20:53:45 2006 +0200
+++ b/src/plugins/quota/quota-private.h	Sun Dec 03 20:55:38 2006 +0200
@@ -14,6 +14,8 @@
 
 	int (*test_alloc)(struct quota_transaction_context *ctx,
 			  uoff_t size, bool *too_large_r);
+
+	unsigned int counting:1;
 };
 
 struct quota_backend_vfuncs {
@@ -81,4 +83,6 @@
 void quota_remove_user_storage(struct quota *quota, 
 			       struct mail_storage *storage);
 
+int quota_count(struct quota *quota, uint64_t *bytes_r, uint64_t *count_r);
+
 #endif

--- a/src/plugins/quota/quota.c	Sun Dec 03 20:53:45 2006 +0200
+++ b/src/plugins/quota/quota.c	Sun Dec 03 20:55:38 2006 +0200
@@ -402,7 +402,12 @@
 	ctx->quota = quota;
 	ctx->box = box;
 	ctx->bytes_left = (uint64_t)-1;
-	ctx->count_left = (uint64_t) -1;
+	ctx->count_left = (uint64_t)-1;
+
+	if (quota->counting) {
+		/* we got here through quota_count_storage() */
+		return ctx;
+	}
 
 	/* find the lowest quota limits from all roots and use them */
 	roots = array_get(&quota->roots, &count);