Repository: relayrides/pushy
Pull request #383 (2016-11-15 05:26:21)

Add a backpressure mechanism for channel writes
* Checks the status of `Channel.isWritable()` before attempting to write a
notification to the channel. If `isWritable()` returns false, return a new failed
future with a cause of `ClientBusyException`. The presence of this exception is
a signal to back off write attempts until the underlying channel has sucessfully returned
to its low water mark
* Adds `ApnsClientBuilder#setChannelWriteBufferWatermark` for setting
 the channel watermark levels
* Adds `ApnsClientBuilder#setByteBufferAllocator` for setting the buffer
 allocator of the netty channel


Patch:
diff --git a/pushy/src/main/java/com/relayrides/pushy/apns/ApnsClient.java b/pushy/src/main/java/com/relayrides/pushy/apns/ApnsClient.java
index 02d8f288..0779a077 100644
--- a/pushy/src/main/java/com/relayrides/pushy/apns/ApnsClient.java
+++ b/pushy/src/main/java/com/relayrides/pushy/apns/ApnsClient.java
@@ -26,6 +26,8 @@ import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.channel.WriteBufferWaterMark;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -160,6 +162,7 @@ public class ApnsClient {
     public static final int ALTERNATE_APNS_PORT = 2197;
 
     private static final ClientNotConnectedException NOT_CONNECTED_EXCEPTION = new ClientNotConnectedException();
+    private static final ClientBusyException CLIENT_BUSY_EXCEPTION = new ClientBusyException();
 
     private static final long INITIAL_RECONNECT_DELAY_SECONDS = 1; // second
     private static final long MAX_RECONNECT_DELAY_SECONDS = 60; // seconds
@@ -265,6 +268,32 @@ public class ApnsClient {
         }
     }
 
+    /**
+     * Sets the write buffer watermark option on the underlying Netty channel
+     *
+     * @param writeBufferWaterMark the write buffer watermark range for the netty channel
+     *
+     * @since 0.9
+     */
+    protected void setChannelWriteBufferWatermark(WriteBufferWaterMark writeBufferWaterMark) {
+        synchronized (this.bootstrap) {
+            this.bootstrap.option(ChannelOption.WRITE_BUFFER_WATER_MARK, writeBufferWaterMark);
+        }
+    }
+
+    /**
+     * Sets the byte buffer allocator on the underlying netty channel
+     *
+     * @param allocator the byte buffer allocator for the netty channel
+     *
+     * @since 0.9
+     */
+    protected void setBufferAllocator(ByteBufAllocator allocator) {
+        synchronized (this.bootstrap) {
+            this.bootstrap.option(ChannelOption.ALLOCATOR, allocator);
+        }
+    }
+
     /**
      * Sets the metrics listener for this client. Metrics listeners gather information that describes the performance
      * and behavior of a client, and are completely optional.
@@ -571,6 +600,9 @@ public class ApnsClient {
         final ChannelPromise connectionReadyPromise = this.connectionReadyPromise;
 
         if (connectionReadyPromise != null && connectionReadyPromise.isSuccess() && connectionReadyPromise.channel().isActive()) {
+            if (!connectionReadyPromise.channel().isWritable()) {
+                return new FailedFuture<>(GlobalEventExecutor.INSTANCE, CLIENT_BUSY_EXCEPTION);
+            }
             final Promise<PushNotificationResponse<ApnsPushNotification>> responsePromise =
                     new DefaultPromise<>(connectionReadyPromise.channel().eventLoop());
 
diff --git a/pushy/src/main/java/com/relayrides/pushy/apns/ApnsClientBuilder.java b/pushy/src/main/java/com/relayrides/pushy/apns/ApnsClientBuilder.java
index 57d73c0f..de7ad0d1 100644
--- a/pushy/src/main/java/com/relayrides/pushy/apns/ApnsClientBuilder.java
+++ b/pushy/src/main/java/com/relayrides/pushy/apns/ApnsClientBuilder.java
@@ -34,6 +34,8 @@ import java.util.concurrent.TimeUnit;
 
 import javax.net.ssl.SSLException;
 
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.channel.WriteBufferWaterMark;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -86,6 +88,9 @@ public class ApnsClientBuilder {
     private Long gracefulShutdownTimeout;
     private TimeUnit gracefulShutdownTimeoutUnit;
 
+    private ByteBufAllocator byteBufferAllocator;
+    private WriteBufferWaterMark channelWriteBufferWaterMark;
+
     private static final Logger log = LoggerFactory.getLogger(ApnsClientBuilder.class);
 
     /**
@@ -348,6 +353,34 @@ public class ApnsClientBuilder {
         return this;
     }
 
+    /**
+     * Sets the buffer allocator for the underlying netty channel
+     *
+     * @param byteBufferAllocator the buffer allocator to use for the underlying netty channel
+     *
+     * @return a reference to this builder
+     *
+     * @since 0.9
+     */
+    public ApnsClientBuilder setByteBufferAllocator(ByteBufAllocator byteBufferAllocator) {
+        this.byteBufferAllocator = byteBufferAllocator;
+        return this;
+    }
+
+    /**
+     * Sets the buffer usage watermark range for the underlying netty channel
+     *
+     * @param writeBufferWatermark the watermark range used in the underlying netty channel
+     *
+     * @return a reference to this builder
+     *
+     * @since 0.9
+     */
+    public ApnsClientBuilder setChannelWriteBufferWatermark(WriteBufferWaterMark writeBufferWatermark) {
+        this.channelWriteBufferWaterMark = writeBufferWatermark;
+        return this;
+    }
+
     /**
      * Constructs a new {@link ApnsClient} with the previously-set configuration.
      *
@@ -416,6 +449,14 @@ public class ApnsClientBuilder {
             apnsClient.setGracefulShutdownTimeout(this.gracefulShutdownTimeoutUnit.toMillis(this.gracefulShutdownTimeout));
         }
 
+        if (this.byteBufferAllocator != null) {
+            apnsClient.setBufferAllocator(byteBufferAllocator);
+        }
+
+        if (this.channelWriteBufferWaterMark != null) {
+            apnsClient.setChannelWriteBufferWatermark(channelWriteBufferWaterMark);
+        }
+
         return apnsClient;
     }
 }
diff --git a/pushy/src/main/java/com/relayrides/pushy/apns/ClientBusyException.java b/pushy/src/main/java/com/relayrides/pushy/apns/ClientBusyException.java
new file mode 100644
index 00000000..9d9020a1
--- /dev/null
+++ b/pushy/src/main/java/com/relayrides/pushy/apns/ClientBusyException.java
@@ -0,0 +1,49 @@
+/* Copyright (c) 2013-2016 RelayRides
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE. */
+
+package com.relayrides.pushy.apns;
+
+/**
+ * An exception thrown to indicate that a notification could not be sent because the client was busy. This is intended
+ * to signal backpressure from the channel associated with this client
+ *
+ * @author <a href="https://github.com/stepheno">Oskar Stephens</a>
+ *
+ * @since 0.9
+ */
+public class ClientBusyException extends IllegalStateException {
+    private static final long serialVersionUID = 1L;
+
+    /**
+     * Constructs a new exception with no message.
+     */
+    public ClientBusyException() {
+        super();
+    }
+
+    /**
+     * Constructs a new exception with the given message.
+     *
+     * @param message a short, human-readable explanation of the cause of this exception
+     */
+    public ClientBusyException(final String message) {
+        super(message);
+    }
+}


Test patch:
diff --git a/pushy/src/test/java/com/relayrides/pushy/apns/ApnsClientTest.java b/pushy/src/test/java/com/relayrides/pushy/apns/ApnsClientTest.java
index 0407d703..de67467b 100644
--- a/pushy/src/test/java/com/relayrides/pushy/apns/ApnsClientTest.java
+++ b/pushy/src/test/java/com/relayrides/pushy/apns/ApnsClientTest.java
@@ -17,6 +17,7 @@ import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import io.netty.channel.WriteBufferWaterMark;
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
@@ -598,6 +599,31 @@ public class ApnsClientTest {
         assertEquals(now, response.getTokenInvalidationTimestamp());
     }
 
+    @Test
+    public void testSendNotificationOnBusyChannel() throws Exception {
+        ApnsClient busyClient;
+
+        try (final InputStream p12InputStream = ApnsClientTest.class.getResourceAsStream(SINGLE_TOPIC_CLIENT_KEYSTORE_FILENAME)) {
+            busyClient = new ApnsClientBuilder()
+                    .setClientCredentials(p12InputStream, KEYSTORE_PASSWORD)
+                    .setTrustedServerCertificateChain(CA_CERTIFICATE)
+                    .setEventLoopGroup(EVENT_LOOP_GROUP)
+                    .setChannelWriteBufferWatermark(new WriteBufferWaterMark(0,0))
+                    .build();
+        }
+        busyClient.connect(HOST, PORT).await();
+
+        final String testToken = ApnsClientTest.generateRandomToken();
+
+        this.server.addToken(DEFAULT_TOPIC, testToken, null);
+
+        final SimpleApnsPushNotification pushNotification = new SimpleApnsPushNotification(testToken, DEFAULT_TOPIC, "test-payload");
+
+        Future<PushNotificationResponse<SimpleApnsPushNotification>> responseFuture = busyClient.sendNotification(pushNotification);
+
+        assertTrue(responseFuture.cause() instanceof ClientBusyException);
+    }
+
     @Test
     public void testSendNotificationWithInternalServerError() throws Exception {
         // Shut down the "normal" server to free the port