1 files changed, 86 insertions, 25 deletions

diff --git a/src/org/gnunet/util/Client.java b/src/org/gnunet/util/Client.java
index d2e1308..fbd94a6 100644
--- a/src/org/gnunet/util/Client.java
+++ b/src/org/gnunet/util/Client.java
@@ -22,6 +22,7 @@ package org.gnunet.util;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+
 /**
  * A connection to a gnunet service.
  *
@@ -69,6 +70,12 @@ public class Client {
     private boolean notifyTransmitReadyDelayed;
 
     /**
+     * When notifyTransmitReadyDelayed is true, this can be used to cancel the task
+     * waiting for the connection to be established.
+     */
+    private Cancelable delayedNotifyTransmitHandle;
+
+    /**
      * Create a connection to a service.
      *
      * @param serviceName name of the service
@@ -111,16 +118,6 @@ public class Client {
         return connection.receive(timeout, receiver);
     }
 
-
-    private static class DelayedTransmitHandle implements Cancelable {
-        Cancelable realTransmitHandle;
-        Cancelable timeoutHandle;
-        @Override
-        public void cancel() {
-            throw new UnsupportedOperationException();
-        }
-    }
-
     /**
      * Ask the client to call us once it is able to send a message.
      *
@@ -132,39 +129,103 @@ public class Client {
      *                    if the caller does not care about temporary connection errors,
      *                    for example because the protocol is stateless
      * @param size        size of the message we want to transmit, can be an upper bound
-     *@param transmitter the MessageTransmitter object to call once the client is ready to transmit or
+     * @param transmitter the MessageTransmitter object to call once the client is ready to transmit or
      *                    when the timeout is over. Guaranteed to be called *after* notifyTransmitReady has returned.  @return a handle that can be used to cancel the transmit request
+     *
+     * @return a handle to cancel the notification
      */
     public Cancelable notifyTransmitReady(final RelativeTime timeout,
                                           final boolean autoRetry, int size, final MessageTransmitter transmitter) {
+        if (notifyTransmitReadyDelayed) {
+            throw new AssertionError("notifyTransmitReady called twice!");
+        }
         if (connection.isConnected()) {
             return connection.notifyTransmitReady(0, timeout, transmitter);
         } else {
             notifyTransmitReadyDelayed = true;
-            final DelayedTransmitHandle delayedTransmitHandle = new DelayedTransmitHandle();
-            delayedTransmitHandle.timeoutHandle = Scheduler.addDelayed(connectBackoff, new Scheduler.Task() {
+            final AbsoluteTime deadline = timeout.toAbsolute();
+            delayedNotifyTransmitHandle = connection.notifyConnected(connectBackoff, new Continuation() {
                 @Override
-                public void run(Scheduler.RunContext ctx) {
-                    if (connection == null) {
-                        return;
-                    }
-                    if (connection.isConnected()) {
+                public void cont(boolean success) {
+                    delayedNotifyTransmitHandle = null;
+                    if (success) {
                         notifyTransmitReadyDelayed = false;
-                        connection.notifyTransmitReady(0, timeout, transmitter);
+                        delayedNotifyTransmitHandle = connection.notifyTransmitReady(0, timeout, new MessageTransmitter() {
+                            @Override
+                            public void transmit(Connection.MessageSink sink) {
+                                delayedNotifyTransmitHandle = null;
+                                transmitter.transmit(sink);
+                            }
+
+                            @Override
+                            public void handleError() {
+                                delayedNotifyTransmitHandle = null;
+                                transmitter.handleError();
+                            }
+                        });
                     } else {
-                        logger.debug("still not connected, retrying in {}ms", connectBackoff.getMilliseconds());
-                        reconnect();
-                        connectBackoff = RelativeTime.min(connectBackoff.multiply(2), MAX_BACKOFF);
-                        Scheduler.addDelayed(connectBackoff, this);
+                        if (deadline.isDue()) {
+                            transmitter.handleError();
+                        } else {
+                            RelativeTime timeout = deadline.getRemaining();
+                            connectBackoff = RelativeTime.min(timeout, RelativeTime.min(connectBackoff.multiply(2), MAX_BACKOFF));
+                            reconnect();
+                            delayedNotifyTransmitHandle = connection.notifyConnected(connectBackoff, this);
+                        }
                     }
                 }
             });
-            return delayedTransmitHandle;
+            return new Cancelable() {
+                @Override
+                public void cancel() {
+                    if (delayedNotifyTransmitHandle != null) {
+                        delayedNotifyTransmitHandle.cancel();
+                    }
+                }
+            };
         }
     }
 
+    /**
+     * Convenience method for sending messages.
+     *
+     * @param timeout when should we give up sending the message, and call cont.cont(false)
+     * @param message the message to send
+     * @param cont called when the message has been sent successfully or on error
+     * @return a handle to cancel sending the message
+     */
+    public Cancelable transmitWhenReady(final RelativeTime timeout, final GnunetMessage.Body message, final Continuation cont) {
+        return notifyTransmitReady(timeout, false, 0, new MessageTransmitter() {
+            @Override
+            public void transmit(Connection.MessageSink sink) {
+                sink.send(message);
+                if (cont != null) {
+                    cont.cont(true);
+                }
+            }
+
+            @Override
+            public void handleError() {
+                if (cont != null) {
+                    cont.cont(false);
+                }
+            }
+        });
+    }
+
+    /**
+     * Convenience method for sending messages. Timeout defaults to FOREVER.
+     *
+     * @param message the message to send
+     * @param cont called when the message has been sent successfully or on error
+     * @return a handle to cancel sending the message
+     */
+    public Cancelable transmitWhenReady(final GnunetMessage.Body message, final Continuation cont) {
+        return transmitWhenReady(RelativeTime.FOREVER, message, cont);
+    }
+
 
-    public void reconnect() {
+    public final void reconnect() {
         if (connection != null) {
             connection.disconnect();
         }