From e4bb5bf5a667e4d7810e0118763661ae9dd889dc Mon Sep 17 00:00:00 2001
From: mjentsch <mtthsjntsch@gmail.com>
Date: Fri, 22 Aug 2014 17:21:46 +0200
Subject: [PATCH] Count pending queries using the elements in the query tree,
 instead of counting the incoming and outgoing messages

Since the server is able to send queries without any response, and is also able to send multiple responses to an arbitrary query, the current way to count pending queries is not accurate and will lead to errors. Instead, the amount of queries in the query tree will be determined to decide whether there are pending messages.
---
 mtproto-client.c | 11 ++++-------
 queries.c        | 10 ++++++----
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/mtproto-client.c b/mtproto-client.c
index 2b8f84a..80bd734 100644
--- a/mtproto-client.c
+++ b/mtproto-client.c
@@ -215,8 +215,6 @@ int rpc_send_packet (struct connection *c) {
 
   self->total_packets_sent ++;
   self->total_data_sent += total_len;
-  self->queries_num ++;
-  logprintf("queries_num=%d\n", self->queries_num);
   return 1;
 }
 
@@ -238,9 +236,6 @@ int rpc_send_message (struct connection *c, void *data, int len) {
 
   self->total_packets_sent ++;
   self->total_data_sent += total_len;
-
-  self->queries_num ++;
-  logprintf("queries_num=%d\n", self->queries_num);
   return 1;
 }
 
@@ -1711,8 +1706,6 @@ int rpc_execute (struct connection *c, int op, int len) {
   logprintf ("outbound rpc connection #%d : received rpc answer %d with %d content bytes\n", c->fd, op, len);
   struct mtproto_connection *self = c->mtconnection;
   
-  self->queries_num --;
-  logprintf ("queries_num=%d\n", c->mtconnection->queries_num);
   /*  
   if (op < 0) {
     assert (read_in (c, Response, Response_len) == Response_len);
@@ -1747,6 +1740,8 @@ int rpc_execute (struct connection *c, int op, int len) {
     return 0;
   case st_client_dh_sent:
     process_auth_complete (c, Response/* + 8*/, Response_len/* - 12*/);
+    self->queries_num --;
+    logprintf ("queries_num=%d\n", c->mtconnection->queries_num);
     if (self->on_ready) {
       self->on_ready(self, self->on_ready_data);
     }
@@ -1786,6 +1781,8 @@ int tc_becomes_ready (struct connection *c) {
   }
   switch (o) {
   case st_init:
+    c->mtconnection->queries_num ++;
+    logprintf ("queries_num=%d\n", c->mtconnection->queries_num);
     send_req_pq_packet (c);
     break;
   case st_authorized:
diff --git a/queries.c b/queries.c
index 86c3170..2b7e845 100644
--- a/queries.c
+++ b/queries.c
@@ -158,6 +158,7 @@ struct query *send_query (struct dc *DC, int ints, void *data, struct query_meth
     }
   }
   queries_tree = tree_insert_query (queries_tree, q, lrand48 ());
+  logprintf("queries_num: %d\n", ++ mtc->queries_num);
 
   q->ev.alarm = (void *)alarm_query;
   q->ev.timeout = get_double_time () + QUERY_TIMEOUT;
@@ -197,6 +198,8 @@ void query_error (long long id) {
       remove_event_timer (&q->ev);
     }
     queries_tree = tree_delete_query (queries_tree, q);
+    logprintf("queries_num: %d\n", -- mtp->queries_num);
+
     if (q->methods && q->methods->on_error) {
       q->methods->on_error (q, error_code, error_len, error);
     } else {
@@ -204,10 +207,9 @@ void query_error (long long id) {
     }
     tfree (q->data, q->data_len * 4);
     tfree (q, sizeof (*q));
+    return;
   }
 
-  mtp->queries_num --;
-  logprintf("queries_num: %d\n", mtp->queries_num);
 }
 
 #define MAX_PACKED_SIZE (1 << 24)
@@ -253,6 +255,8 @@ void query_result (long long id UU) {
       remove_event_timer (&q->ev);
     }
     queries_tree = tree_delete_query (queries_tree, q);
+    logprintf("queries_num: %d\n", -- mtp->queries_num);
+
     if (q->methods && q->methods->on_answer) {
       q->methods->on_answer (q);
       assert (mtp->in_ptr == mtp->in_end);
@@ -264,8 +268,6 @@ void query_result (long long id UU) {
     mtp->in_ptr = end;
     mtp->in_end = eend;
   }
-  //queries_num --;
-  //logprintf("queries_num: %d\n", queries_num);
 } 
 
 #define event_timer_cmp(a,b) ((a)->timeout > (b)->timeout ? 1 : ((a)->timeout < (b)->timeout ? -1 : (memcmp (a, b, sizeof (struct event_timer)))))