1 files changed, 22 insertions, 10 deletions

diff --git a/src/service/peerstore/perf_peerstore_store.c b/src/service/peerstore/perf_peerstore_store.c
index e59af61e5..e328be93e 100644
--- a/src/service/peerstore/perf_peerstore_store.c
+++ b/src/service/peerstore/perf_peerstore_store.c
@@ -32,16 +32,20 @@ static int ok = 1;
 
 static struct GNUNET_PEERSTORE_Handle *h;
 
+static struct GNUNET_PEERSTORE_WatchContext *wc;
+
 static char *ss = "test_peerstore_stress";
 static struct GNUNET_PeerIdentity p;
 static char *k = "test_peerstore_stress_key";
 static char *v = "test_peerstore_stress_val";
 
 static int count = 0;
+static int count_fin = 0;
 
 static void
-disconnect ()
+disconnect (void *cls)
 {
+  GNUNET_PEERSTORE_watch_cancel (wc);
   if (NULL != h)
     GNUNET_PEERSTORE_disconnect (h);
   GNUNET_SCHEDULER_shutdown ();
@@ -49,14 +53,27 @@ disconnect ()
 
 
 static void
+store_cont (void *cls, int ret)
+{
+  count_fin++;
+  if (count_fin == count)
+  {
+    ok = 0;
+    GNUNET_SCHEDULER_add_now (&disconnect, NULL);
+  }
+}
+
+
+static void
 store ()
 {
+  count++;
   GNUNET_PEERSTORE_store (h, ss, &p, k, v, strlen (v) + 1,
                           GNUNET_TIME_UNIT_FOREVER_ABS,
                           (count ==
                            0) ? GNUNET_PEERSTORE_STOREOPTION_REPLACE :
-                          GNUNET_PEERSTORE_STOREOPTION_MULTIPLE, NULL, NULL);
-  count++;
+                          GNUNET_PEERSTORE_STOREOPTION_MULTIPLE, store_cont,
+                          NULL);
 }
 
 
@@ -65,12 +82,7 @@ watch_cb (void *cls, const struct GNUNET_PEERSTORE_Record *record,
           const char *emsg)
 {
   GNUNET_assert (NULL == emsg);
-  if (STORES == count)
-  {
-    ok = 0;
-    disconnect ();
-  }
-  else
+  if (STORES > count)
     store ();
 }
 
@@ -82,7 +94,7 @@ run (void *cls, const struct GNUNET_CONFIGURATION_Handle *cfg,
   memset (&p, 5, sizeof(p));
   h = GNUNET_PEERSTORE_connect (cfg);
   GNUNET_assert (NULL != h);
-  GNUNET_PEERSTORE_watch (h, ss, &p, k, &watch_cb, NULL);
+  wc = GNUNET_PEERSTORE_watch (h, ss, &p, k, &watch_cb, NULL);
   store ();
 }