mirror of
https://github.com/signalapp/Signal-Server
synced 2026-04-26 01:08:05 +01:00
Use a custom redis pubsub implementation rather than Jedis.
// FREEBIE
This commit is contained in:
Moxie Marlinspike
@@ -0,0 +1,7 @@
|
||||
package org.whispersystems.dispatch;
|
||||
|
||||
public interface DispatchChannel {
|
||||
public void onDispatchMessage(String channel, byte[] message);
|
||||
public void onDispatchSubscribed(String channel);
|
||||
public void onDispatchUnsubscribed(String channel);
|
||||
}
|
||||
191
src/main/java/org/whispersystems/dispatch/DispatchManager.java
Normal file
191
src/main/java/org/whispersystems/dispatch/DispatchManager.java
Normal file
@@ -0,0 +1,191 @@
|
||||
package org.whispersystems.dispatch;
|
||||
|
||||
import com.google.common.base.Optional;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.whispersystems.dispatch.io.RedisPubSubConnectionFactory;
|
||||
import org.whispersystems.dispatch.redis.PubSubConnection;
|
||||
import org.whispersystems.dispatch.redis.PubSubReply;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.Collection;
|
||||
import java.util.HashMap;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.Executor;
|
||||
import java.util.concurrent.Executors;
|
||||
|
||||
public class DispatchManager extends Thread {
|
||||
|
||||
private final Logger logger = LoggerFactory.getLogger(DispatchManager.class);
|
||||
private final Executor executor = Executors.newCachedThreadPool();
|
||||
private final Map<String, DispatchChannel> subscriptions = new HashMap<>();
|
||||
|
||||
private final Optional<DispatchChannel> deadLetterChannel;
|
||||
private final RedisPubSubConnectionFactory redisPubSubConnectionFactory;
|
||||
|
||||
private PubSubConnection pubSubConnection;
|
||||
private volatile boolean running;
|
||||
|
||||
public DispatchManager(RedisPubSubConnectionFactory redisPubSubConnectionFactory,
|
||||
Optional<DispatchChannel> deadLetterChannel)
|
||||
{
|
||||
this.redisPubSubConnectionFactory = redisPubSubConnectionFactory;
|
||||
this.deadLetterChannel = deadLetterChannel;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void start() {
|
||||
this.pubSubConnection = redisPubSubConnectionFactory.connect();
|
||||
this.running = true;
|
||||
super.start();
|
||||
}
|
||||
|
||||
public void shutdown() {
|
||||
this.running = false;
|
||||
this.pubSubConnection.close();
|
||||
}
|
||||
|
||||
public void subscribe(String name, DispatchChannel dispatchChannel) {
|
||||
Optional<DispatchChannel> previous;
|
||||
|
||||
synchronized (subscriptions) {
|
||||
previous = Optional.fromNullable(subscriptions.get(name));
|
||||
subscriptions.put(name, dispatchChannel);
|
||||
}
|
||||
|
||||
try {
|
||||
pubSubConnection.subscribe(name);
|
||||
} catch (IOException e) {
|
||||
logger.warn("Subscription error", e);
|
||||
}
|
||||
|
||||
if (previous.isPresent()) {
|
||||
dispatchUnsubscription(name, previous.get());
|
||||
}
|
||||
}
|
||||
|
||||
public void unsubscribe(String name, DispatchChannel channel) {
|
||||
final Optional<DispatchChannel> subscription;
|
||||
|
||||
synchronized (subscriptions) {
|
||||
subscription = Optional.fromNullable(subscriptions.get(name));
|
||||
|
||||
if (subscription.isPresent() && subscription.get() == channel) {
|
||||
subscriptions.remove(name);
|
||||
}
|
||||
}
|
||||
|
||||
if (subscription.isPresent()) {
|
||||
try {
|
||||
pubSubConnection.unsubscribe(name);
|
||||
} catch (IOException e) {
|
||||
logger.warn("Unsubscribe error", e);
|
||||
}
|
||||
|
||||
dispatchUnsubscription(name, subscription.get());
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public void run() {
|
||||
while (running) {
|
||||
try {
|
||||
PubSubReply reply = pubSubConnection.read();
|
||||
|
||||
switch (reply.getType()) {
|
||||
case UNSUBSCRIBE: break;
|
||||
case SUBSCRIBE: dispatchSubscribe(reply); break;
|
||||
case MESSAGE: dispatchMessage(reply); break;
|
||||
default: throw new AssertionError("Unknown pubsub reply type! " + reply.getType());
|
||||
}
|
||||
} catch (IOException e) {
|
||||
logger.warn("***** PubSub Connection Error *****", e);
|
||||
if (running) {
|
||||
this.pubSubConnection.close();
|
||||
this.pubSubConnection = redisPubSubConnectionFactory.connect();
|
||||
resubscribe();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.warn("DispatchManager Shutting Down...");
|
||||
}
|
||||
|
||||
private void dispatchSubscribe(final PubSubReply reply) {
|
||||
final Optional<DispatchChannel> subscription;
|
||||
|
||||
synchronized (subscriptions) {
|
||||
subscription = Optional.fromNullable(subscriptions.get(reply.getChannel()));
|
||||
}
|
||||
|
||||
if (subscription.isPresent()) {
|
||||
dispatchSubscription(reply.getChannel(), subscription.get());
|
||||
} else {
|
||||
logger.warn("Received subscribe event for non-existing channel: " + reply.getChannel());
|
||||
}
|
||||
}
|
||||
|
||||
private void dispatchMessage(PubSubReply reply) {
|
||||
Optional<DispatchChannel> subscription;
|
||||
|
||||
synchronized (subscriptions) {
|
||||
subscription = Optional.fromNullable(subscriptions.get(reply.getChannel()));
|
||||
}
|
||||
|
||||
if (subscription.isPresent()) {
|
||||
dispatchMessage(reply.getChannel(), subscription.get(), reply.getContent().get());
|
||||
} else if (deadLetterChannel.isPresent()) {
|
||||
dispatchMessage(reply.getChannel(), deadLetterChannel.get(), reply.getContent().get());
|
||||
} else {
|
||||
logger.warn("Received message for non-existing channel, with no dead letter handler: " + reply.getChannel());
|
||||
}
|
||||
}
|
||||
|
||||
private void resubscribe() {
|
||||
final Collection<String> names;
|
||||
|
||||
synchronized (subscriptions) {
|
||||
names = subscriptions.keySet();
|
||||
}
|
||||
|
||||
new Thread() {
|
||||
@Override
|
||||
public void run() {
|
||||
try {
|
||||
for (String name : names) {
|
||||
pubSubConnection.subscribe(name);
|
||||
}
|
||||
} catch (IOException e) {
|
||||
logger.warn("***** RESUBSCRIPTION ERROR *****", e);
|
||||
}
|
||||
}
|
||||
}.start();
|
||||
}
|
||||
|
||||
private void dispatchMessage(final String name, final DispatchChannel channel, final byte[] message) {
|
||||
executor.execute(new Runnable() {
|
||||
@Override
|
||||
public void run() {
|
||||
channel.onDispatchMessage(name, message);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
private void dispatchSubscription(final String name, final DispatchChannel channel) {
|
||||
executor.execute(new Runnable() {
|
||||
@Override
|
||||
public void run() {
|
||||
channel.onDispatchSubscribed(name);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
private void dispatchUnsubscription(final String name, final DispatchChannel channel) {
|
||||
executor.execute(new Runnable() {
|
||||
@Override
|
||||
public void run() {
|
||||
channel.onDispatchUnsubscribed(name);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package org.whispersystems.dispatch.io;
|
||||
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
|
||||
public class RedisInputStream {
|
||||
|
||||
private static final byte CR = 0x0D;
|
||||
private static final byte LF = 0x0A;
|
||||
|
||||
private final InputStream inputStream;
|
||||
|
||||
public RedisInputStream(InputStream inputStream) {
|
||||
this.inputStream = inputStream;
|
||||
}
|
||||
|
||||
public String readLine() throws IOException {
|
||||
ByteArrayOutputStream boas = new ByteArrayOutputStream();
|
||||
|
||||
boolean foundCr = false;
|
||||
|
||||
while (true) {
|
||||
int character = inputStream.read();
|
||||
|
||||
if (character == -1) {
|
||||
throw new IOException("Stream closed!");
|
||||
}
|
||||
|
||||
boas.write(character);
|
||||
|
||||
if (foundCr && character == LF) break;
|
||||
else if (character == CR) foundCr = true;
|
||||
else if (foundCr) foundCr = false;
|
||||
}
|
||||
|
||||
byte[] data = boas.toByteArray();
|
||||
return new String(data, 0, data.length-2);
|
||||
}
|
||||
|
||||
public byte[] readFully(int size) throws IOException {
|
||||
byte[] result = new byte[size];
|
||||
int offset = 0;
|
||||
int remaining = result.length;
|
||||
|
||||
while (remaining > 0) {
|
||||
int read = inputStream.read(result, offset, remaining);
|
||||
|
||||
if (read < 0) {
|
||||
throw new IOException("Stream closed!");
|
||||
}
|
||||
|
||||
offset += read;
|
||||
remaining -= read;
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
public void close() throws IOException {
|
||||
inputStream.close();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
package org.whispersystems.dispatch.io;
|
||||
|
||||
import org.whispersystems.dispatch.redis.PubSubConnection;
|
||||
|
||||
public interface RedisPubSubConnectionFactory {
|
||||
|
||||
public PubSubConnection connect();
|
||||
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
package org.whispersystems.dispatch.redis;
|
||||
|
||||
import com.google.common.base.Optional;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.whispersystems.dispatch.io.RedisInputStream;
|
||||
import org.whispersystems.dispatch.redis.protocol.ArrayReplyHeader;
|
||||
import org.whispersystems.dispatch.redis.protocol.IntReply;
|
||||
import org.whispersystems.dispatch.redis.protocol.StringReplyHeader;
|
||||
import org.whispersystems.dispatch.util.Util;
|
||||
|
||||
import java.io.BufferedInputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.OutputStream;
|
||||
import java.net.Socket;
|
||||
import java.util.Arrays;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
|
||||
public class PubSubConnection {
|
||||
|
||||
private final Logger logger = LoggerFactory.getLogger(PubSubConnection.class);
|
||||
|
||||
private static final byte[] UNSUBSCRIBE_TYPE = {'u', 'n', 's', 'u', 'b', 's', 'c', 'r', 'i', 'b', 'e' };
|
||||
private static final byte[] SUBSCRIBE_TYPE = {'s', 'u', 'b', 's', 'c', 'r', 'i', 'b', 'e' };
|
||||
private static final byte[] MESSAGE_TYPE = {'m', 'e', 's', 's', 'a', 'g', 'e' };
|
||||
|
||||
private static final byte[] SUBSCRIBE_COMMAND = {'S', 'U', 'B', 'S', 'C', 'R', 'I', 'B', 'E', ' ' };
|
||||
private static final byte[] UNSUBSCRIBE_COMMAND = {'U', 'N', 'S', 'U', 'B', 'S', 'C', 'R', 'I', 'B', 'E', ' '};
|
||||
private static final byte[] CRLF = {'\r', '\n' };
|
||||
|
||||
private final OutputStream outputStream;
|
||||
private final RedisInputStream inputStream;
|
||||
private final Socket socket;
|
||||
private final AtomicBoolean closed;
|
||||
|
||||
public PubSubConnection(Socket socket) throws IOException {
|
||||
this.socket = socket;
|
||||
this.outputStream = socket.getOutputStream();
|
||||
this.inputStream = new RedisInputStream(new BufferedInputStream(socket.getInputStream()));
|
||||
this.closed = new AtomicBoolean(false);
|
||||
}
|
||||
|
||||
public void subscribe(String channelName) throws IOException {
|
||||
if (closed.get()) throw new IOException("Connection closed!");
|
||||
|
||||
byte[] command = Util.combine(SUBSCRIBE_COMMAND, channelName.getBytes(), CRLF);
|
||||
outputStream.write(command);
|
||||
}
|
||||
|
||||
public void unsubscribe(String channelName) throws IOException {
|
||||
if (closed.get()) throw new IOException("Connection closed!");
|
||||
|
||||
byte[] command = Util.combine(UNSUBSCRIBE_COMMAND, channelName.getBytes(), CRLF);
|
||||
outputStream.write(command);
|
||||
}
|
||||
|
||||
public PubSubReply read() throws IOException {
|
||||
if (closed.get()) throw new IOException("Connection closed!");
|
||||
|
||||
ArrayReplyHeader replyHeader = new ArrayReplyHeader(inputStream.readLine());
|
||||
|
||||
if (replyHeader.getElementCount() != 3) {
|
||||
throw new IOException("Received array reply header with strange count: " + replyHeader.getElementCount());
|
||||
}
|
||||
|
||||
StringReplyHeader replyTypeHeader = new StringReplyHeader(inputStream.readLine());
|
||||
byte[] replyType = inputStream.readFully(replyTypeHeader.getStringLength());
|
||||
inputStream.readLine();
|
||||
|
||||
if (Arrays.equals(SUBSCRIBE_TYPE, replyType)) return readSubscribeReply();
|
||||
else if (Arrays.equals(UNSUBSCRIBE_TYPE, replyType)) return readUnsubscribeReply();
|
||||
else if (Arrays.equals(MESSAGE_TYPE, replyType)) return readMessageReply();
|
||||
else throw new IOException("Unknown reply type: " + new String(replyType));
|
||||
}
|
||||
|
||||
public void close() {
|
||||
try {
|
||||
this.closed.set(true);
|
||||
this.inputStream.close();
|
||||
this.outputStream.close();
|
||||
this.socket.close();
|
||||
} catch (IOException e) {
|
||||
logger.warn("Exception while closing", e);
|
||||
}
|
||||
}
|
||||
|
||||
private PubSubReply readMessageReply() throws IOException {
|
||||
StringReplyHeader channelNameHeader = new StringReplyHeader(inputStream.readLine());
|
||||
byte[] channelName = inputStream.readFully(channelNameHeader.getStringLength());
|
||||
inputStream.readLine();
|
||||
|
||||
StringReplyHeader messageHeader = new StringReplyHeader(inputStream.readLine());
|
||||
byte[] message = inputStream.readFully(messageHeader.getStringLength());
|
||||
inputStream.readLine();
|
||||
|
||||
return new PubSubReply(PubSubReply.Type.MESSAGE, new String(channelName), Optional.of(message));
|
||||
}
|
||||
|
||||
private PubSubReply readUnsubscribeReply() throws IOException {
|
||||
String channelName = readSubscriptionReply();
|
||||
return new PubSubReply(PubSubReply.Type.UNSUBSCRIBE, channelName, Optional.<byte[]>absent());
|
||||
}
|
||||
|
||||
private PubSubReply readSubscribeReply() throws IOException {
|
||||
String channelName = readSubscriptionReply();
|
||||
return new PubSubReply(PubSubReply.Type.SUBSCRIBE, channelName, Optional.<byte[]>absent());
|
||||
}
|
||||
|
||||
private String readSubscriptionReply() throws IOException {
|
||||
StringReplyHeader channelNameHeader = new StringReplyHeader(inputStream.readLine());
|
||||
byte[] channelName = inputStream.readFully(channelNameHeader.getStringLength());
|
||||
inputStream.readLine();
|
||||
|
||||
IntReply subscriptionCount = new IntReply(inputStream.readLine());
|
||||
|
||||
return new String(channelName);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package org.whispersystems.dispatch.redis;
|
||||
|
||||
import com.google.common.base.Optional;
|
||||
|
||||
public class PubSubReply {
|
||||
|
||||
public enum Type {
|
||||
MESSAGE,
|
||||
SUBSCRIBE,
|
||||
UNSUBSCRIBE
|
||||
}
|
||||
|
||||
private final Type type;
|
||||
private final String channel;
|
||||
private final Optional<byte[]> content;
|
||||
|
||||
public PubSubReply(Type type, String channel, Optional<byte[]> content) {
|
||||
this.type = type;
|
||||
this.channel = channel;
|
||||
this.content = content;
|
||||
}
|
||||
|
||||
public Type getType() {
|
||||
return type;
|
||||
}
|
||||
|
||||
public String getChannel() {
|
||||
return channel;
|
||||
}
|
||||
|
||||
public Optional<byte[]> getContent() {
|
||||
return content;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package org.whispersystems.dispatch.redis.protocol;
|
||||
|
||||
import java.io.IOException;
|
||||
|
||||
public class ArrayReplyHeader {
|
||||
|
||||
private final int elementCount;
|
||||
|
||||
public ArrayReplyHeader(String header) throws IOException {
|
||||
if (header == null || header.length() < 2 || header.charAt(0) != '*') {
|
||||
throw new IOException("Invalid array reply header: " + header);
|
||||
}
|
||||
|
||||
try {
|
||||
this.elementCount = Integer.parseInt(header.substring(1));
|
||||
} catch (NumberFormatException e) {
|
||||
throw new IOException(e);
|
||||
}
|
||||
}
|
||||
|
||||
public int getElementCount() {
|
||||
return elementCount;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package org.whispersystems.dispatch.redis.protocol;
|
||||
|
||||
import java.io.IOException;
|
||||
|
||||
public class IntReply {
|
||||
|
||||
private final int value;
|
||||
|
||||
public IntReply(String reply) throws IOException {
|
||||
if (reply == null || reply.length() < 2 || reply.charAt(0) != ':') {
|
||||
throw new IOException("Invalid int reply: " + reply);
|
||||
}
|
||||
|
||||
try {
|
||||
this.value = Integer.parseInt(reply.substring(1));
|
||||
} catch (NumberFormatException e) {
|
||||
throw new IOException(e);
|
||||
}
|
||||
}
|
||||
|
||||
public int getValue() {
|
||||
return value;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package org.whispersystems.dispatch.redis.protocol;
|
||||
|
||||
import java.io.IOException;
|
||||
|
||||
public class StringReplyHeader {
|
||||
|
||||
private final int stringLength;
|
||||
|
||||
public StringReplyHeader(String header) throws IOException {
|
||||
if (header == null || header.length() < 2 || header.charAt(0) != '$') {
|
||||
throw new IOException("Invalid string reply header: " + header);
|
||||
}
|
||||
|
||||
try {
|
||||
this.stringLength = Integer.parseInt(header.substring(1));
|
||||
} catch (NumberFormatException e) {
|
||||
throw new IOException(e);
|
||||
}
|
||||
}
|
||||
|
||||
public int getStringLength() {
|
||||
return stringLength;
|
||||
}
|
||||
}
|
||||
36
src/main/java/org/whispersystems/dispatch/util/Util.java
Normal file
36
src/main/java/org/whispersystems/dispatch/util/Util.java
Normal file
@@ -0,0 +1,36 @@
|
||||
package org.whispersystems.dispatch.util;
|
||||
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.io.IOException;
|
||||
|
||||
public class Util {
|
||||
|
||||
public static byte[] combine(byte[]... elements) {
|
||||
try {
|
||||
int sum = 0;
|
||||
|
||||
for (byte[] element : elements) {
|
||||
sum += element.length;
|
||||
}
|
||||
|
||||
ByteArrayOutputStream baos = new ByteArrayOutputStream(sum);
|
||||
|
||||
for (byte[] element : elements) {
|
||||
baos.write(element);
|
||||
}
|
||||
|
||||
return baos.toByteArray();
|
||||
} catch (IOException e) {
|
||||
throw new AssertionError(e);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
public static void sleep(long millis) {
|
||||
try {
|
||||
Thread.sleep(millis);
|
||||
} catch (InterruptedException e) {
|
||||
throw new AssertionError(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user