dict: Support large iterations by sending data in output stream flush callback.

Fixes also a hang if a lot of changes were done during the iteration.

--HG--
branch : HEAD

diff --git a/src/dict/dict-server.c b/src/dict/dict-server.c
index 2c70476f355dced9542f2064786c961297ecb5a5..ca2bd56d5dc6ef99add50db5ce21086dc8cd3b28 100644 (file)
--- a/src/dict/dict-server.c
+++ b/src/dict/dict-server.c
@@ -6,6 +6,7 @@
 #include "network.h"
 #include "istream.h"
 #include "ostream.h"
+#include "str.h"
 #include "dict.h"
 #include "dict-client.h"
 #include "dict-server.h"
@@ -13,6 +14,8 @@
 #include <stdlib.h>
 #include <unistd.h>
 
+#define DICT_OUTPUT_OPTIMAL_SIZE 1024
+
 struct dict_server_transaction {
        unsigned int id;
        struct dict_transaction_context *ctx;
@@ -32,6 +35,8 @@ struct dict_client_connection {
        struct istream *input;
        struct ostream *output;
 
+       struct dict_iterate_context *iter_ctx;
+
        /* There are only a few transactions per client, so keeping them in
           array is fast enough */
        ARRAY_DEFINE(transactions, struct dict_server_transaction);
@@ -58,6 +63,11 @@ static int cmd_lookup(struct dict_client_connection *conn, const char *line)
        const char *value;
        int ret;
 
+       if (conn->iter_ctx != NULL) {
+               i_error("dict client: LOOKUP: Can't lookup while iterating");
+               return -1;
+       }
+
        /* <key> */
        ret = dict_lookup(conn->dict, pool_datastack_create(), line, &value);
        if (ret > 0) {
@@ -73,13 +83,46 @@ static int cmd_lookup(struct dict_client_connection *conn, const char *line)
        return 0;
 }
 
-static int cmd_iterate(struct dict_client_connection *conn, const char *line)
+static int cmd_iterate_flush(struct dict_client_connection *conn)
 {
-       struct dict_iterate_context *ctx;
-       const char *const *args;
+       string_t *str;
        const char *key, *value;
        int ret;
 
+       str = t_str_new(256);
+       o_stream_cork(conn->output);
+       while ((ret = dict_iterate(conn->iter_ctx, &key, &value)) > 0) {
+               str_truncate(str, 0);
+               str_printfa(str, "%s\t%s\n", key, value);
+               o_stream_send(conn->output, str_data(str), str_len(str));
+
+               if (o_stream_get_buffer_used_size(conn->output) >
+                   DICT_OUTPUT_OPTIMAL_SIZE) {
+                       if (o_stream_flush(conn->output) <= 0)
+                               break;
+                       /* flushed everything, continue */
+               }
+       }
+
+       if (ret <= 0) {
+               /* finished iterating */
+               o_stream_unset_flush_callback(conn->output);
+               dict_iterate_deinit(&conn->iter_ctx);
+               o_stream_send(conn->output, "\n", 1);
+       }
+       o_stream_uncork(conn->output);
+       return ret <= 0 ? 1 : 0;
+}
+
+static int cmd_iterate(struct dict_client_connection *conn, const char *line)
+{
+       const char *const *args;
+
+       if (conn->iter_ctx != NULL) {
+               i_error("dict client: ITERATE: Already iterating");
+               return -1;
+       }
+
        args = t_strsplit(line, "\t");
        if (str_array_length(args) != 2) {
                i_error("dict client: ITERATE: broken input");
@@ -87,22 +130,10 @@ static int cmd_iterate(struct dict_client_connection *conn, const char *line)
        }
 
        /* <flags> <path> */
-       o_stream_cork(conn->output);
-       ctx = dict_iterate_init(conn->dict, args[1], atoi(args[0]));
-       while ((ret = dict_iterate(ctx, &key, &value)) > 0) {
-               /* FIXME: we don't want to keep blocking here. set a flush
-                  function and send the replies there when buffer gets full */
-               T_BEGIN {
-                       const char *reply;
-
-                       reply = t_strdup_printf("%s\t%s\n", key, value);
-                       o_stream_send_str(conn->output, reply);
-               } T_END;
-       }
-       dict_iterate_deinit(&ctx);
+       conn->iter_ctx = dict_iterate_init(conn->dict, args[1], atoi(args[0]));
 
-       o_stream_send_str(conn->output, "\n");
-       o_stream_uncork(conn->output);
+       o_stream_set_flush_callback(conn->output, cmd_iterate_flush, conn);
+       cmd_iterate_flush(conn);
        return 0;
 }
 
@@ -193,6 +224,11 @@ static int cmd_commit(struct dict_client_connection *conn, const char *line)
        const char *reply;
        int ret;
 
+       if (conn->iter_ctx != NULL) {
+               i_error("dict client: COMMIT: Can't commit while iterating");
+               return -1;
+       }
+
        if (dict_server_transaction_lookup_parse(conn, line, &trans) < 0)
                return -1;
 
@@ -431,6 +467,9 @@ static void dict_client_connection_deinit(struct dict_client_connection *conn)
                array_free(&conn->transactions);
        }
 
+       if (conn->iter_ctx != NULL)
+               dict_iterate_deinit(&conn->iter_ctx);
+
        io_remove(&conn->io);
        i_stream_destroy(&conn->input);
        o_stream_destroy(&conn->output);