/**
* Copyright 2016 Yahoo Inc.
*
* 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.yahoo.pulsar.websocket;
import static com.yahoo.pulsar.websocket.WebSocketError.FailedToCreateProducer;
import static com.yahoo.pulsar.websocket.WebSocketError.FailedToDeserializeFromJSON;
import static com.yahoo.pulsar.websocket.WebSocketError.PayloadEncodingError;
import static com.yahoo.pulsar.websocket.WebSocketError.UnknownError;
import static java.lang.String.format;
import java.io.IOException;
import java.util.Base64;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLongFieldUpdater;
import java.util.concurrent.atomic.LongAdder;
import javax.servlet.http.HttpServletRequest;
import org.eclipse.jetty.websocket.api.Session;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.yahoo.pulsar.client.api.CompressionType;
import com.yahoo.pulsar.client.api.Message;
import com.yahoo.pulsar.client.api.MessageBuilder;
import com.yahoo.pulsar.client.api.Producer;
import com.yahoo.pulsar.client.api.ProducerConfiguration;
import com.yahoo.pulsar.client.api.ProducerConfiguration.MessageRoutingMode;
import com.yahoo.pulsar.common.naming.DestinationName;
import com.yahoo.pulsar.common.util.ObjectMapperFactory;
import com.yahoo.pulsar.websocket.data.ProducerAck;
import com.yahoo.pulsar.websocket.data.ProducerMessage;
import com.yahoo.pulsar.websocket.stats.StatsBuckets;
/**
* Websocket end-point url handler to handle incoming message coming from client. Websocket end-point url handler to
* handle incoming message coming from client.
* <p>
* On every produced message from client it calls broker to persists it.
* </p>
*
*/
public class ProducerHandler extends AbstractWebSocketHandler {
private Producer producer;
private final LongAdder numMsgsSent;
private final LongAdder numMsgsFailed;
private final LongAdder numBytesSent;
private final StatsBuckets publishLatencyStatsUSec;
private volatile long msgPublishedCounter = 0;
private static final AtomicLongFieldUpdater<ProducerHandler> MSG_PUBLISHED_COUNTER_UPDATER =
AtomicLongFieldUpdater.newUpdater(ProducerHandler.class, "msgPublishedCounter");
public static final long[] ENTRY_LATENCY_BUCKETS_USEC = { 500, 1_000, 5_000, 10_000, 20_000, 50_000, 100_000,
200_000, 1000_000 };
public ProducerHandler(WebSocketService service, HttpServletRequest request) {
super(service, request);
this.numMsgsSent = new LongAdder();
this.numBytesSent = new LongAdder();
this.numMsgsFailed = new LongAdder();
this.publishLatencyStatsUSec = new StatsBuckets(ENTRY_LATENCY_BUCKETS_USEC);
}
@Override
public void close() throws IOException {
if (producer != null) {
this.service.removeProducer(this);
producer.close();
}
}
@Override
public void onWebSocketConnect(Session session) {
super.onWebSocketConnect(session);
try {
ProducerConfiguration conf = getProducerConfiguration();
this.producer = service.getPulsarClient().createProducer(topic, conf);
this.service.addProducer(this);
} catch (Exception e) {
close(FailedToCreateProducer, e.getMessage());
}
}
@Override
public void onWebSocketText(String message) {
ProducerMessage sendRequest;
byte[] rawPayload = null;
String requestContext = null;
try {
sendRequest = ObjectMapperFactory.getThreadLocal().readValue(message, ProducerMessage.class);
requestContext = sendRequest.context;
rawPayload = Base64.getDecoder().decode(sendRequest.payload);
} catch (IOException e) {
sendAckResponse(new ProducerAck(FailedToDeserializeFromJSON, e.getMessage(), null, null));
return;
} catch (IllegalArgumentException e) {
String msg = format("Invalid Base64 message-payload error=%s", e.getMessage());
sendAckResponse(new ProducerAck(PayloadEncodingError, msg, null, requestContext));
return;
}
final long msgSize = rawPayload.length;
MessageBuilder builder = MessageBuilder.create().setContent(rawPayload);
if (sendRequest.properties != null) {
builder.setProperties(sendRequest.properties);
}
if (sendRequest.key != null) {
builder.setKey(sendRequest.key);
}
if (sendRequest.replicationClusters != null) {
builder.setReplicationClusters(sendRequest.replicationClusters);
}
Message msg = builder.build();
final long now = System.nanoTime();
producer.sendAsync(msg).thenAccept(msgId -> {
updateSentMsgStats(msgSize, TimeUnit.NANOSECONDS.toMicros(System.nanoTime() - now));
if (isConnected()) {
String messageId = Base64.getEncoder().encodeToString(msgId.toByteArray());
sendAckResponse(new ProducerAck(messageId, sendRequest.context));
}
}).exceptionally(exception -> {
numMsgsFailed.increment();
sendAckResponse(
new ProducerAck(UnknownError, exception.getMessage(), null, sendRequest.context));
return null;
});
}
public Producer getProducer() {
return this.producer;
}
public long getAndResetNumMsgsSent() {
return numMsgsSent.sumThenReset();
}
public long getAndResetNumBytesSent() {
return numBytesSent.sumThenReset();
}
public long getAndResetNumMsgsFailed() {
return numMsgsFailed.sumThenReset();
}
public long[] getAndResetPublishLatencyStatsUSec() {
publishLatencyStatsUSec.refresh();
return publishLatencyStatsUSec.getBuckets();
}
public StatsBuckets getPublishLatencyStatsUSec() {
return this.publishLatencyStatsUSec;
}
public long getMsgPublishedCounter() {
return MSG_PUBLISHED_COUNTER_UPDATER.get(this);
}
protected CompletableFuture<Boolean> isAuthorized(String authRole) {
return service.getAuthorizationManager().canProduceAsync(DestinationName.get(topic), authRole);
}
private void sendAckResponse(ProducerAck response) {
try {
String msg = ObjectMapperFactory.getThreadLocal().writeValueAsString(response);
getSession().getRemote().sendString(msg);
} catch (JsonProcessingException e) {
log.warn("[{}] Failed to generate ack json-response {}", producer.getTopic(), e.getMessage(), e);
} catch (Exception e) {
log.warn("[{}] Failed to send ack {}", producer.getTopic(), e.getMessage(), e);
}
}
private void updateSentMsgStats(long msgSize, long latencyUsec) {
this.publishLatencyStatsUSec.addValue(latencyUsec);
this.numBytesSent.add(msgSize);
this.numMsgsSent.increment();
MSG_PUBLISHED_COUNTER_UPDATER.getAndIncrement(this);
}
private ProducerConfiguration getProducerConfiguration() {
ProducerConfiguration conf = new ProducerConfiguration();
// Set to false to prevent the server thread from being blocked if a lot of messages are pending.
conf.setBlockIfQueueFull(false);
if (queryParams.containsKey("sendTimeoutMillis")) {
conf.setSendTimeout(Integer.parseInt(queryParams.get("sendTimeoutMillis")), TimeUnit.MILLISECONDS);
}
if (queryParams.containsKey("batchingEnabled")) {
conf.setBatchingEnabled(Boolean.parseBoolean(queryParams.get("batchingEnabled")));
}
if (queryParams.containsKey("batchingMaxMessages")) {
conf.setBatchingMaxMessages(Integer.parseInt(queryParams.get("batchingMaxMessages")));
}
if (queryParams.containsKey("maxPendingMessages")) {
conf.setMaxPendingMessages(Integer.parseInt(queryParams.get("maxPendingMessages")));
}
if (queryParams.containsKey("batchingMaxPublishDelay")) {
conf.setBatchingMaxPublishDelay(Integer.parseInt(queryParams.get("batchingMaxPublishDelay")),
TimeUnit.MILLISECONDS);
}
if (queryParams.containsKey("messageRoutingMode")) {
conf.setMessageRoutingMode(MessageRoutingMode.valueOf(queryParams.get("messageRoutingMode")));
}
if (queryParams.containsKey("compressionType")) {
conf.setCompressionType(CompressionType.valueOf(queryParams.get("compressionType")));
}
return conf;
}
private static final Logger log = LoggerFactory.getLogger(ProducerHandler.class);
}