/*
* Copyright (c) 2011-2015 Spotify AB
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.spotify.google.cloud.pubsub.client;
import com.google.common.collect.ImmutableSet;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;
import java.util.List;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
import static com.google.common.collect.Iterables.concat;
import static com.spotify.google.cloud.pubsub.client.AssertWithTimeout.assertThatWithin;
import static java.util.Arrays.asList;
import static java.util.Collections.singletonList;
import static java.util.concurrent.TimeUnit.DAYS;
import static java.util.concurrent.TimeUnit.SECONDS;
import static java.util.stream.Collectors.toList;
import static java.util.stream.IntStream.range;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.notNullValue;
import static org.junit.Assert.assertThat;
import static org.mockito.Matchers.anyListOf;
import static org.mockito.Matchers.anyString;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@RunWith(MockitoJUnitRunner.class)
public class PublisherTest {
@Mock Pubsub pubsub;
@Mock Publisher.Listener listener;
@Captor ArgumentCaptor<PubsubFuture<List<String>>> batchFutureCaptor;
private final ConcurrentMap<String, BlockingQueue<Request>> topics = new ConcurrentHashMap<>();
private Publisher publisher;
@Before
public void setUp() {
setUpPubsubClient();
publisher = Publisher.builder()
.project("test")
.pubsub(pubsub)
.listener(listener)
.build();
}
@After
public void tearDown() throws Exception {
publisher.close();
}
@Test
public void testConfigurationGetters() {
final Publisher publisher = Publisher.builder()
.pubsub(pubsub)
.project("test")
.concurrency(11)
.batchSize(12)
.queueSize(13)
.build();
assertThat(publisher.project(), is("test"));
assertThat(publisher.concurrency(), is(11));
assertThat(publisher.batchSize(), is(12));
assertThat(publisher.queueSize(), is(13));
}
@Test
public void testOutstandingRequests() throws InterruptedException, ExecutionException {
final LinkedBlockingQueue<Request> t1 = new LinkedBlockingQueue<>();
final LinkedBlockingQueue<Request> t2 = new LinkedBlockingQueue<>();
topics.put("t1", t1);
topics.put("t2", t2);
final Message m1 = Message.of("1");
final Message m2 = Message.of("2");
// Verify that the outstanding requests before publishing anything is 0
assertThat(publisher.outstandingRequests(), is(0));
// Publish a message and verify that the outstanding request counter rises to 1
final CompletableFuture<String> f1 = publisher.publish("t1", m1);
assertThatWithin(5, SECONDS, publisher::outstandingRequests, is(1));
// Publish another message and verify that the outstanding request counter rises to 2
final CompletableFuture<String> f2 = publisher.publish("t2", m2);
assertThatWithin(5, SECONDS, publisher::outstandingRequests, is(2));
// Respond to the first request and verify that the outstanding request counter falls to 1
t1.take().future.succeed(singletonList("id1"));
assertThatWithin(5, SECONDS, publisher::outstandingRequests, is(1));
// Respond to the second request and verify that the outstanding request counter falls to 0
t2.take().future.succeed(singletonList("id2"));
assertThatWithin(5, SECONDS, publisher::outstandingRequests, is(0));
}
@Test
public void testLatencyBoundedBatchingSingleMessage() throws InterruptedException, ExecutionException {
final LinkedBlockingQueue<Request> t = new LinkedBlockingQueue<>();
topics.put("t", t);
final Message m = Message.of("1");
// Publish two messages
publisher.publish("t", m);
// Check that the publisher eventually times out gathering messages for the batch and sends the single message.
final Request request = t.take();
assertThat(request.messages.size(), is(1));
}
@Test
public void testLatencyBoundedBatchingTwoMessages() throws InterruptedException, ExecutionException {
final LinkedBlockingQueue<Request> t = new LinkedBlockingQueue<>();
topics.put("t", t);
final Message m1 = Message.of("1");
final Message m2 = Message.of("2");
// Publish two messages
publisher.publish("t", m1);
publisher.publish("t", m2);
// Check that the publisher eventually times out gathering messages for the batch and sends just the two messages.
final Request request = t.take();
assertThat(request.messages.size(), is(2));
}
@Test
public void testSizeBoundedBatching() throws InterruptedException, ExecutionException {
final LinkedBlockingQueue<Request> t = new LinkedBlockingQueue<>();
topics.put("t", t);
publisher = Publisher.builder()
.project("test")
.pubsub(pubsub)
.batchSize(2)
.maxLatencyMs(DAYS.toMillis(1))
.build();
final Message m1 = Message.of("1");
final Message m2 = Message.of("2");
// Publish a single message
publisher.publish("t", m1);
// Verify that the batch is not sent
Thread.sleep(1000);
verify(pubsub, never()).publish(anyString(), anyString(), anyListOf(Message.class));
// Send one more message, completing the batch.
publisher.publish("t", m2);
// Check that the batch got sent.
verify(pubsub, timeout(5000)).publish(anyString(), anyString(), anyListOf(Message.class));
final Request request = t.take();
assertThat(request.messages.size(), is(2));
}
@Test
public void testPendingTopics() throws InterruptedException, ExecutionException {
final LinkedBlockingQueue<Request> t1 = new LinkedBlockingQueue<>();
final LinkedBlockingQueue<Request> t2 = new LinkedBlockingQueue<>();
topics.put("t1", t1);
topics.put("t2", t2);
publisher = Publisher.builder()
.project("test")
.pubsub(pubsub)
.listener(listener)
.concurrency(1)
.maxLatencyMs(1)
.build();
final Message m1 = Message.of("1");
final Message m2 = Message.of("2");
// Verify that the pending topics before publishing anything is 0
assertThat(publisher.pendingTopics(), is(0));
// Publish a message and verify that the pending topics is still 0
final CompletableFuture<String> f1 = publisher.publish("t1", m1);
verify(pubsub, timeout(1000)).publish("test", "t1", singletonList(m1));
assertThat(publisher.pendingTopics(), is(0));
// Publish a message on a different topic and verify that the pending topics counter rises to 1
publisher.publish("t2", m2);
assertThatWithin(5, SECONDS, publisher::pendingTopics, is(1));
// Respond to the first request and verify that the pending topics falls to 0
t1.take().future.succeed(singletonList("id1"));
assertThatWithin(5, SECONDS, publisher::pendingTopics, is(0));
}
@Test
public void testListener() throws InterruptedException, ExecutionException {
setUpPubsubClient();
publisher = Publisher.builder()
.project("test")
.pubsub(pubsub)
.listener(listener)
.concurrency(1)
.build();
final LinkedBlockingQueue<Request> t1 = new LinkedBlockingQueue<>();
final LinkedBlockingQueue<Request> t2 = new LinkedBlockingQueue<>();
topics.put("t1", t1);
topics.put("t2", t2);
final Message m1 = Message.of("1");
final Message m2a = Message.of("2a");
final Message m2b = Message.of("2b");
// Verify that the listener got called when the publisher was created
verify(listener, timeout(5000)).publisherCreated(publisher);
// Publish a message and verify that the listener got called
final CompletableFuture<String> f1 = publisher.publish("t1", m1);
verify(listener, timeout(5000)).publishingMessage(publisher, "t1", m1, f1);
verify(listener, timeout(5000)).sendingBatch(
eq(publisher), eq("t1"), eq(singletonList(m1)), batchFutureCaptor.capture());
// Publish two messages on a different topic and verify that the listener got told that the topic is pending
final CompletableFuture<String> f2a = publisher.publish("t2", m2a);
final CompletableFuture<String> f2b = publisher.publish("t2", m2b);
verify(listener, timeout(5000)).publishingMessage(publisher, "t2", m2a, f2a);
verify(listener, timeout(5000)).publishingMessage(publisher, "t2", m2b, f2b);
verify(listener, timeout(5000)).topicPending(publisher, "t2", 1, 1);
// Respond to the first request and verify that the batch future is completed
t1.take().future.succeed(singletonList("id1"));
final List<String> batchIds1 = batchFutureCaptor.getValue().get();
assertThat(batchIds1, contains("id1"));
// verify that the listener got called for the second batch
verify(listener, timeout(5000)).sendingBatch(
eq(publisher), eq("t2"), eq(asList(m2a, m2b)), batchFutureCaptor.capture());
// Respond to the second requests and verify that the batch future is completed
t2.take().future.succeed(asList("id2a", "id2b"));
final List<String> batchIds2 = batchFutureCaptor.getValue().get();
assertThat(batchIds2, contains("id2a", "id2b"));
// Close the publisher and verify that the listener got called
publisher.close();
verify(listener, timeout(5000)).publisherClosed(publisher);
}
@Test
public void testListenerAdapter() throws Exception {
final CompletableFuture<Void> created = new CompletableFuture<>();
final CompletableFuture<Void> closed = new CompletableFuture<>();
final Publisher.ListenerAdapter listener = new Publisher.ListenerAdapter() {
@Override
public void publisherCreated(final Publisher publisher) {
created.complete(null);
}
@Override
public void publisherClosed(final Publisher publisher) {
closed.complete(null);
}
};
publisher = Publisher.builder()
.project("test")
.pubsub(pubsub)
.listener(listener)
.concurrency(1)
.build();
assertThat(created.isDone(), is(true));
publisher.close();
publisher.closeFuture().get();
assertThat(closed.isDone(), is(true));
}
@Test
public void testQueueSize() throws InterruptedException, ExecutionException {
final LinkedBlockingQueue<Request> t = new LinkedBlockingQueue<>();
topics.put("t", t);
final Message m0 = Message.of("0");
final Message m1 = Message.of("1");
final Message m2 = Message.of("2");
publisher = Publisher.builder()
.project("test")
.pubsub(pubsub)
.listener(listener)
.concurrency(1)
.queueSize(1)
.build();
// Send one message to saturate the concurrency (queue size == 0)
final CompletableFuture<String> f0 = publisher.publish("t", m0);
verify(pubsub, timeout(1000)).publish(eq("test"), eq("t"), eq(singletonList(m0)));
// Send one message to occupy the queue (queue size == 1)
final CompletableFuture<String> f1 = publisher.publish("t", m1);
// Send a message that should fast fail due to the queue being full
final CompletableFuture<String> f2 = publisher.publish("t", m2);
assertThat(f2.isCompletedExceptionally(), is(true));
assertThat(exception(f2), is(instanceOf(QueueFullException.class)));
// Complete the first request
t.take().future.succeed(singletonList("id0"));
// Verify that the second one is sent and complete it
verify(pubsub, timeout(1000)).publish(eq("test"), eq("t"), eq(singletonList(m1)));
t.take().future.succeed(singletonList("id1"));
// Verify that the fast-failed request was not sent
Thread.sleep(1000);
verify(pubsub, never()).publish(anyString(), anyString(), eq(singletonList(m2)));
}
@Test
public void verifyConcurrentBacklogConsumption() throws Exception {
final LinkedBlockingQueue<Request> t = new LinkedBlockingQueue<>();
topics.put("t", t);
publisher = Publisher.builder()
.project("test")
.pubsub(pubsub)
.listener(listener)
.concurrency(2)
.batchSize(2)
.queueSize(100)
.build();
// Saturate concurrency with two messages
final Message m0a = Message.of("0a");
final CompletableFuture<String> f0a = publisher.publish("t", m0a);
final Request r0a = t.take();
final Message m0b = Message.of("0b");
final CompletableFuture<String> f0b = publisher.publish("t", m0b);
final Request r0b = t.take();
// Enqueue enough for at least two more batches
final List<Message> m1 = range(0, 4).mapToObj(String::valueOf).map(Message::of).collect(toList());
final List<CompletableFuture<String>> f1 = m1.stream().map(m -> publisher.publish("t", m)).collect(toList());
// Complete the first two requests
r0a.future.succeed(singletonList("0a"));
r0b.future.succeed(singletonList("0b"));
// Verify that two batches kicked off concurrently and that we got all four messages in the two batches
final Request r1a = t.poll(30, SECONDS);
final Request r1b = t.poll(30, SECONDS);
assertThat(r1a, is(notNullValue()));
assertThat(r1b, is(notNullValue()));
final Set<Message> r1received = ImmutableSet.copyOf(concat(r1a.messages, r1b.messages));
assertThat(r1received, is(ImmutableSet.copyOf(m1)));
}
private Throwable exception(final CompletableFuture<?> f) {
if (!f.isCompletedExceptionally()) {
throw new IllegalArgumentException();
}
try {
f.get();
} catch (InterruptedException e) {
return e;
} catch (ExecutionException e) {
return e.getCause();
}
return null;
}
private void setUpPubsubClient() {
reset(pubsub);
when(pubsub.publish(anyString(), anyString(), anyListOf(Message.class)))
.thenAnswer(invocation -> {
final String topic = (String) invocation.getArguments()[1];
@SuppressWarnings("unchecked") final List<Message> messages =
(List<Message>) invocation.getArguments()[2];
final RequestInfo requestInfo = RequestInfo.builder()
.operation("publish")
.method("POST")
.uri("/publish")
.payloadSize(4711)
.build();
final PubsubFuture<List<String>> future = new PubsubFuture<>(requestInfo);
final BlockingQueue<Request> queue = topics.get(topic);
queue.add(new Request(messages, future));
return future;
});
}
private static class Request {
final List<Message> messages;
final PubsubFuture<List<String>> future;
Request(final List<Message> messages,
final PubsubFuture<List<String>> future) {
this.messages = messages;
this.future = future;
}
}
}