11 files changed, 117 insertions, 89 deletions

diff --git a/src/dht/gnunet-dht-get.c b/src/dht/gnunet-dht-get.c
index 6846b3af7..8a1b836f6 100644
--- a/src/dht/gnunet-dht-get.c
+++ b/src/dht/gnunet-dht-get.c
@@ -82,9 +82,14 @@ static unsigned int result_count;
  */
 static int ret;
 
+/**
+ * Task scheduled to handle timeout.
+ */
+static struct GNUNET_SCHEDULER_Task *tt;
+
 
 /**
- * Task run to clean up on timeout.
+ * Task run to clean up on shutdown.
  *
  * @param cls unused
  */
@@ -101,6 +106,24 @@ cleanup_task (void *cls)
     GNUNET_DHT_disconnect (dht_handle);
     dht_handle = NULL;
   }
+  if (NULL != tt)
+  {
+    GNUNET_SCHEDULER_cancel (tt);
+    tt = NULL;
+  }
+}
+
+
+/**
+ * Task run on timeout. Triggers shutdown.
+ *
+ * @param cls unused
+ */
+static void
+timeout_task (void *cls)
+{
+  tt = NULL;
+  GNUNET_SCHEDULER_shutdown ();
 }
 
 
@@ -170,8 +193,9 @@ run (void *cls, char *const *args, const char *cfgfile,
   GNUNET_CRYPTO_hash (query_key, strlen (query_key), &key);
   if (verbose)
     FPRINTF (stderr, "%s `%s' \n",  _("Issueing DHT GET with key"), GNUNET_h2s_full (&key));
-  GNUNET_SCHEDULER_add_delayed (timeout_request,
-                                &cleanup_task, NULL);
+  GNUNET_SCHEDULER_add_shutdown (&cleanup_task, NULL);
+  tt = GNUNET_SCHEDULER_add_delayed (timeout_request,
+                                     &timeout_task, NULL);
   get_handle =
       GNUNET_DHT_get_start (dht_handle, query_type, &key, replication,
                             (demultixplex_everywhere) ? GNUNET_DHT_RO_DEMULTIPLEX_EVERYWHERE : GNUNET_DHT_RO_NONE,
diff --git a/src/dht/gnunet-dht-monitor.c b/src/dht/gnunet-dht-monitor.c
index 0502496d7..7f14255d3 100644
--- a/src/dht/gnunet-dht-monitor.c
+++ b/src/dht/gnunet-dht-monitor.c
@@ -71,6 +71,11 @@ static unsigned int result_count;
  */
 static int ret;
 
+/**
+ * Task scheduled to handle timeout.
+ */
+static struct GNUNET_SCHEDULER_Task *tt;
+
 
 /**
  * Stop monitoring request and start shutdown
@@ -92,9 +97,28 @@ cleanup_task (void *cls)
     GNUNET_DHT_disconnect (dht_handle);
     dht_handle = NULL;
   }
+  if (NULL != tt)
+  {
+    GNUNET_SCHEDULER_cancel (tt);
+    tt = NULL;
+  }
+}
+
+
+/**
+ * We hit a timeout. Stop monitoring request and start shutdown
+ *
+ * @param cls closure (unused)
+ */
+static void
+timeout_task (void *cls)
+{
+  tt = NULL;
+  GNUNET_SCHEDULER_shutdown ();
 }
 
 
+
 /**
  * Callback called on each GET request going through the DHT.
  *
@@ -107,7 +131,7 @@ cleanup_task (void *cls)
  * @param desired_replication_level Desired replication level.
  * @param key Key of the requested data.
  */
-void
+static void
 get_callback (void *cls,
               enum GNUNET_DHT_RouteOption options,
               enum GNUNET_BLOCK_Type type,
@@ -139,7 +163,7 @@ get_callback (void *cls,
  * @param data Pointer to the result data.
  * @param size Number of bytes in data.
  */
-void
+static void
 get_resp_callback (void *cls,
                    enum GNUNET_BLOCK_Type type,
                    const struct GNUNET_PeerIdentity *get_path,
@@ -177,7 +201,7 @@ get_resp_callback (void *cls,
  * @param data Pointer to the data carried.
  * @param size Number of bytes in data.
  */
-void
+static void
 put_callback (void *cls,
               enum GNUNET_DHT_RouteOption options,
               enum GNUNET_BLOCK_Type type,
@@ -242,7 +266,11 @@ run (void *cls, char *const *args, const char *cfgfile,
     FPRINTF (stderr,
              "Monitoring for %s\n",
              GNUNET_STRINGS_relative_time_to_string (timeout_request, GNUNET_NO));
-  GNUNET_SCHEDULER_add_delayed (timeout_request, &cleanup_task, NULL);
+  tt = GNUNET_SCHEDULER_add_delayed (timeout_request,
+                                     &timeout_task,
+                                     NULL);
+  GNUNET_SCHEDULER_add_shutdown (&cleanup_task,
+                                NULL);
   monitor_handle = GNUNET_DHT_monitor_start (dht_handle,
                                              block_type,
                                              key,
diff --git a/src/dht/gnunet-dht-put.c b/src/dht/gnunet-dht-put.c
index 19e385be1..972369e10 100644
--- a/src/dht/gnunet-dht-put.c
+++ b/src/dht/gnunet-dht-put.c
@@ -98,6 +98,7 @@ shutdown_task (void *cls)
   }
 }
 
+
 /**
  * Signature of the main function of a task.
  *
@@ -132,6 +133,7 @@ message_sent_cont (void *cls, int success)
   GNUNET_SCHEDULER_add_now (&shutdown_task, NULL);
 }
 
+
 /**
  * Main function that will be run by the scheduler.
  *
@@ -141,14 +143,15 @@ message_sent_cont (void *cls, int success)
  * @param c configuration
  */
 static void
-run (void *cls, char *const *args, const char *cfgfile,
+run (void *cls,
+     char *const *args,
+     const char *cfgfile,
      const struct GNUNET_CONFIGURATION_Handle *c)
 {
   struct GNUNET_TIME_Relative timeout;
   struct GNUNET_TIME_Absolute expiration;
 
   cfg = c;
-
   if ((NULL == query_key) || (NULL == data))
   {
     FPRINTF (stderr, "%s",  _("Must provide KEY and DATA for DHT put!\n"));
diff --git a/src/dht/gnunet-service-dht.c b/src/dht/gnunet-service-dht.c
index 1aeca9b85..a44be2dfe 100644
--- a/src/dht/gnunet-service-dht.c
+++ b/src/dht/gnunet-service-dht.c
@@ -161,9 +161,8 @@ run (void *cls, struct GNUNET_SERVER_Handle *server,
     shutdown_task (NULL);
     return;
   }
-  GNUNET_SCHEDULER_add_delayed (GNUNET_TIME_UNIT_FOREVER_REL,
-                                &shutdown_task,
-                                NULL);
+  GNUNET_SCHEDULER_add_shutdown (&shutdown_task,
+                                 NULL);
   GDS_transport_handle =
       GNUNET_TRANSPORT_connect (GDS_cfg, NULL, NULL, NULL, NULL, NULL);
   if (GDS_transport_handle == NULL)
diff --git a/src/dht/gnunet-service-dht_clients.c b/src/dht/gnunet-service-dht_clients.c
index e42069fe0..a5757c25a 100644
--- a/src/dht/gnunet-service-dht_clients.c
+++ b/src/dht/gnunet-service-dht_clients.c
@@ -456,12 +456,8 @@ transmit_next_request_task (void *cls)
 {
   struct ClientQueryRecord *cqr;
   struct GNUNET_TIME_Relative delay;
-  const struct GNUNET_SCHEDULER_TaskContext *tc;
 
   retry_task = NULL;
-  tc = GNUNET_SCHEDULER_get_task_context ();
-  if (0 != (tc->reason & GNUNET_SCHEDULER_REASON_SHUTDOWN))
-    return;
   while (NULL != (cqr = GNUNET_CONTAINER_heap_remove_root (retry_heap)))
   {
     cqr->hnode = NULL;
diff --git a/src/dht/gnunet-service-dht_neighbours.c b/src/dht/gnunet-service-dht_neighbours.c
index d82effb34..8ab39c9aa 100644
--- a/src/dht/gnunet-service-dht_neighbours.c
+++ b/src/dht/gnunet-service-dht_neighbours.c
@@ -712,12 +712,8 @@ send_find_peer_message (void *cls)
   struct GNUNET_TIME_Relative next_send_time;
   struct BloomConstructorContext bcc;
   struct GNUNET_CONTAINER_BloomFilter *peer_bf;
-  const struct GNUNET_SCHEDULER_TaskContext *tc;
 
   find_peer_task = NULL;
-  tc = GNUNET_SCHEDULER_get_task_context ();
-  if ((tc->reason & GNUNET_SCHEDULER_REASON_SHUTDOWN) != 0)
-    return;
   if (newly_found_peers > bucket_size)
   {
     /* If we are finding many peers already, no need to send out our request right now! */
@@ -2443,6 +2439,7 @@ GDS_NEIGHBOURS_done ()
   }
 }
 
+
 /**
  * Get the ID of the local node.
  *
diff --git a/src/dht/gnunet-service-wdht.c b/src/dht/gnunet-service-wdht.c
index 0668c9078..b58bb729d 100644
--- a/src/dht/gnunet-service-wdht.c
+++ b/src/dht/gnunet-service-wdht.c
@@ -102,9 +102,8 @@ run (void *cls, struct GNUNET_SERVER_Handle *server,
     shutdown_task (NULL);
     return;
   }
-  GNUNET_SCHEDULER_add_delayed (GNUNET_TIME_UNIT_FOREVER_REL,
-                                &shutdown_task,
-                                NULL);
+  GNUNET_SCHEDULER_add_shutdown (&shutdown_task,
+                                 NULL);
 }
 
 
diff --git a/src/dht/gnunet-service-wdht_clients.c b/src/dht/gnunet-service-wdht_clients.c
index c3aff03aa..8eb114660 100644
--- a/src/dht/gnunet-service-wdht_clients.c
+++ b/src/dht/gnunet-service-wdht_clients.c
@@ -880,12 +880,8 @@ transmit_next_request_task (void *cls)
 {
   struct ClientQueryRecord *cqr;
   struct GNUNET_TIME_Relative delay;
-  const struct GNUNET_SCHEDULER_TaskContext *tc;
 
   retry_task = NULL;
-  tc = GNUNET_SCHEDULER_get_task_context ();
-  if (0 != (tc->reason & GNUNET_SCHEDULER_REASON_SHUTDOWN))
-    return;
   while (NULL != (cqr = GNUNET_CONTAINER_heap_remove_root (retry_heap)))
   {
     cqr->hnode = NULL;
diff --git a/src/dht/gnunet-service-xdht.c b/src/dht/gnunet-service-xdht.c
index f6c7c5167..8a60a8d9d 100644
--- a/src/dht/gnunet-service-xdht.c
+++ b/src/dht/gnunet-service-xdht.c
@@ -17,14 +17,12 @@
      Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
      Boston, MA 02110-1301, USA.
 */
-
 /**
  * @file dht/gnunet-service-xdht.c
  * @brief GNUnet DHT service
  * @author Christian Grothoff
  * @author Nathan Evans
  */
-
 #include "platform.h"
 #include "gnunet_block_lib.h"
 #include "gnunet_util_lib.h"
@@ -125,9 +123,8 @@ run (void *cls, struct GNUNET_SERVER_Handle *server,
     shutdown_task (NULL);
     return;
   }
-  GNUNET_SCHEDULER_add_delayed (GNUNET_TIME_UNIT_FOREVER_REL,
-                                &shutdown_task,
-                                NULL);
+  GNUNET_SCHEDULER_add_shutdown (&shutdown_task,
+                                 NULL);
 }
 
 
diff --git a/src/dht/gnunet-service-xdht_clients.c b/src/dht/gnunet-service-xdht_clients.c
index e08ccf792..3475d00c9 100644
--- a/src/dht/gnunet-service-xdht_clients.c
+++ b/src/dht/gnunet-service-xdht_clients.c
@@ -877,12 +877,8 @@ transmit_next_request_task (void *cls)
 {
   struct ClientQueryRecord *cqr;
   struct GNUNET_TIME_Relative delay;
-  const struct GNUNET_SCHEDULER_TaskContext *tc;
 
   retry_task = NULL;
-  tc = GNUNET_SCHEDULER_get_task_context ();
-  if (0 != (tc->reason & GNUNET_SCHEDULER_REASON_SHUTDOWN))
-    return;
   while (NULL != (cqr = GNUNET_CONTAINER_heap_remove_root (retry_heap)))
   {
     cqr->hnode = NULL;
diff --git a/src/dht/gnunet_dht_profiler.c b/src/dht/gnunet_dht_profiler.c
index f0c4ec7be..03ea62152 100644
--- a/src/dht/gnunet_dht_profiler.c
+++ b/src/dht/gnunet_dht_profiler.c
@@ -461,9 +461,21 @@ do_shutdown (void *cls)
     a_ctx = NULL;
   }
   //FIXME: Should we collect stats only for put/get not for other messages.
-  if(NULL != bandwidth_stats_op)
+  if (NULL != bandwidth_stats_op)
+  {
     GNUNET_TESTBED_operation_done (bandwidth_stats_op);
-  bandwidth_stats_op = NULL;
+    bandwidth_stats_op = NULL;
+  }
+  if (NULL != successor_stats_op)
+  {
+    GNUNET_TESTBED_operation_done (successor_stats_op);
+    successor_stats_op = NULL;
+  }
+  if (NULL != successor_stats_task)
+  {
+    GNUNET_SCHEDULER_cancel (successor_stats_task);
+    successor_stats_task = NULL;
+  }
   GNUNET_free_non_null (a_ac);
 }
 
@@ -687,11 +699,7 @@ teardown_dht_connection (void *cls)
 {
   struct Context *ctx = cls;
   struct GNUNET_TESTBED_Operation *op;
-  const struct GNUNET_SCHEDULER_TaskContext *tc;
 
-  tc = GNUNET_SCHEDULER_get_task_context ();
-  if (0 != (GNUNET_SCHEDULER_REASON_SHUTDOWN & tc->reason))
-    return;
   GNUNET_assert (NULL != ctx);
   GNUNET_assert (NULL != (op = ctx->op));
   ctx->op = NULL;
@@ -1070,7 +1078,7 @@ successor_stats_cont (void *cls,
   successor_stats_op = NULL;
   if (0 == max_searches)
   {
-    start_func();
+    start_func ();
     return;
   }
 
@@ -1108,39 +1116,27 @@ successor_stats_cont (void *cls,
                                                                  GNUNET_NO);
   if ((start_val == val) && (count == num_peers))
   {
-    DEBUG("CIRCLE COMPLETED after %u tries", tries);
+    DEBUG ("CIRCLE COMPLETED after %u tries", tries);
     if(NULL == successor_stats_task)
     {
       start_func();
     }
     return;
   }
-  else
+  if (max_searches == ++tries)
   {
-    if (max_searches == ++tries)
-    {
-      GNUNET_log (GNUNET_ERROR_TYPE_DEBUG,
-                  "Maximum tries %u exceeded while checking successor TOTAL TRIES %u"
-                  " circle formation.  Exiting\n",
-                  max_searches,tries);
-      if (NULL != successor_stats_task)
-      {
-        successor_stats_task = NULL;
-      }
-      if(NULL == successor_stats_task)
-      {
-        start_func();
-      }
-
-      return;
-    }
-    else
-    {
-      flag = 0;
-      successor_stats_task = GNUNET_SCHEDULER_add_delayed (delay_stats,
-                                                           &collect_stats, cls);
-    }
+    GNUNET_log (GNUNET_ERROR_TYPE_DEBUG,
+                "Maximum tries %u exceeded while checking successor TOTAL TRIES %u"
+                " circle formation.  Exiting\n",
+                max_searches,tries);
+    start_func();
+    return;
   }
+  flag = 0;
+  successor_stats_task
+    = GNUNET_SCHEDULER_add_delayed (delay_stats,
+                                    &collect_stats,
+                                    cls);
 }
 
 
@@ -1214,26 +1210,21 @@ successor_stats_iterator (void *cls,
 static void
 collect_stats (void *cls)
 {
-  const struct GNUNET_SCHEDULER_TaskContext *tc;
-
-  tc = GNUNET_SCHEDULER_get_task_context ();
-  if ((GNUNET_SCHEDULER_REASON_SHUTDOWN & tc->reason) != 0)
-    return;
-
+  successor_stats_task = NULL;
   GNUNET_log (GNUNET_ERROR_TYPE_INFO,
               "Start collecting statistics...\n");
   GNUNET_assert(NULL != testbed_handles);
 
   if (0 != max_searches)
-  successor_peer_hashmap = GNUNET_CONTAINER_multihashmap_create (num_peers,
-                                                                    GNUNET_NO);
-  successor_stats_op =
-          GNUNET_TESTBED_get_statistics (num_peers, testbed_handles,
-                                         "dht", NULL,
-                                          successor_stats_iterator,
-                                          successor_stats_cont, cls);
-
-  GNUNET_assert(NULL != successor_stats_op);
+    successor_peer_hashmap
+      = GNUNET_CONTAINER_multihashmap_create (num_peers,
+                                              GNUNET_NO);
+  successor_stats_op
+    = GNUNET_TESTBED_get_statistics (num_peers, testbed_handles,
+                                     "dht", NULL,
+                                     successor_stats_iterator,
+                                     successor_stats_cont, cls);
+  GNUNET_assert (NULL != successor_stats_op);
 }
 
 
@@ -1265,9 +1256,10 @@ service_started (void *cls,
      collect_stat_cls->service_connect_ctx = cls;
      collect_stat_cls->op = op;
 
-     successor_stats_task = GNUNET_SCHEDULER_add_delayed (delay_stats,
-                                                          &collect_stats,
-                                                          collect_stat_cls);
+     successor_stats_task
+       = GNUNET_SCHEDULER_add_delayed (delay_stats,
+                                       &collect_stats,
+                                       collect_stat_cls);
   }
 }
 
@@ -1393,7 +1385,8 @@ run (void *cls, char *const *args, const char *cfgfile,
 
   if (0 == num_peers)
   {
-    GNUNET_log (GNUNET_ERROR_TYPE_ERROR, _("Exiting as the number of peers is %u\n"),
+    GNUNET_log (GNUNET_ERROR_TYPE_ERROR,
+                _("Exiting as the number of peers is %u\n"),
                 num_peers);
     return;
   }
@@ -1401,8 +1394,8 @@ run (void *cls, char *const *args, const char *cfgfile,
   event_mask = 0;
   GNUNET_TESTBED_run (hosts_file, cfg, num_peers, event_mask, NULL,
                       NULL, &test_run, NULL);
-  GNUNET_SCHEDULER_add_delayed (GNUNET_TIME_UNIT_FOREVER_REL, &do_shutdown,
-                                NULL);
+  GNUNET_SCHEDULER_add_shutdown (&do_shutdown,
+                                 NULL);
 }