server/task/task-distributed/src/main/java/org/apache/james/task/eventsourcing/distributed/RabbitMQWorkQueue.java - james-project - Git at Google

 /**
  * *************************************************************
  * Licensed to the Apache Software Foundation (ASF) under one   *
  * or more contributor license agreements.  See the NOTICE file *
  * distributed with this work for additional information        *
  * regarding copyright ownership.  The ASF licenses this file   *
  * to you 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 org.apache.james.task.eventsourcing.distributed;

 import static com.rabbitmq.client.MessageProperties.PERSISTENT_TEXT_PLAIN;
 import static org.apache.james.backends.rabbitmq.Constants.AUTO_DELETE;
 import static org.apache.james.backends.rabbitmq.Constants.DURABLE;
 import static org.apache.james.backends.rabbitmq.Constants.REQUEUE;

 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Optional;

 import org.apache.james.backends.rabbitmq.RabbitMQConfiguration;
 import org.apache.james.backends.rabbitmq.ReceiverProvider;
 import org.apache.james.server.task.json.JsonTaskSerializer;
 import org.apache.james.task.Task;
 import org.apache.james.task.TaskId;
 import org.apache.james.task.TaskManagerWorker;
 import org.apache.james.task.TaskWithId;
 import org.apache.james.task.WorkQueue;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.rabbitmq.client.AMQP;
 import com.rabbitmq.client.Delivery;

 import reactor.core.Disposable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.UnicastProcessor;
 import reactor.core.scheduler.Schedulers;
 import reactor.rabbitmq.AcknowledgableDelivery;
 import reactor.rabbitmq.BindingSpecification;
 import reactor.rabbitmq.ConsumeOptions;
 import reactor.rabbitmq.ExchangeSpecification;
 import reactor.rabbitmq.OutboundMessage;
 import reactor.rabbitmq.QueueSpecification;
 import reactor.rabbitmq.Receiver;
 import reactor.rabbitmq.Sender;
 import reactor.util.retry.Retry;

 public class RabbitMQWorkQueue implements WorkQueue {
     private static final Logger LOGGER = LoggerFactory.getLogger(RabbitMQWorkQueue.class);

     static final String EXCHANGE_NAME = "taskManagerWorkQueueExchange";
     static final String QUEUE_NAME = "taskManagerWorkQueue";
     static final String ROUTING_KEY = "taskManagerWorkQueueRoutingKey";

     static final String CANCEL_REQUESTS_EXCHANGE_NAME = "taskManagerCancelRequestsExchange";
     static final String CANCEL_REQUESTS_ROUTING_KEY = "taskManagerCancelRequestsRoutingKey";
     public static final String TASK_ID = "taskId";

     public static final int NUM_RETRIES = 8;
     public static final Duration FIRST_BACKOFF = Duration.ofMillis(100);
     public static final boolean AUTO_DELETE_QUEUE = true;

     private final TaskManagerWorker worker;
     private final JsonTaskSerializer taskSerializer;
     private final RabbitMQWorkQueueConfiguration configuration;
     private final Sender sender;
     private final ReceiverProvider receiverProvider;
     private final RabbitMQConfiguration rabbitMQConfiguration;
     private final CancelRequestQueueName cancelRequestQueueName;
     private UnicastProcessor<TaskId> sendCancelRequestsQueue;
     private Disposable sendCancelRequestsQueueHandle;
     private Disposable receiverHandle;
     private Disposable cancelRequestListenerHandle;

     public RabbitMQWorkQueue(TaskManagerWorker worker, Sender sender,
                              ReceiverProvider receiverProvider, JsonTaskSerializer taskSerializer,
                              RabbitMQWorkQueueConfiguration configuration, CancelRequestQueueName cancelRequestQueueName,
                              RabbitMQConfiguration rabbitMQConfiguration) {
         this.cancelRequestQueueName = cancelRequestQueueName;
         this.worker = worker;
         this.receiverProvider = receiverProvider;
         this.sender = sender;
         this.taskSerializer = taskSerializer;
         this.configuration = configuration;
         this.rabbitMQConfiguration = rabbitMQConfiguration;
     }

     @Override
     public void start() {
         startWorkqueue();
         listenToCancelRequests();
     }

     private void startWorkqueue() {
         declareQueue();

         if (configuration.enabled()) {
             consumeWorkqueue();
         }
     }

     @VisibleForTesting
     void declareQueue() {
         Mono<AMQP.Exchange.DeclareOk> declareExchange = sender
             .declareExchange(ExchangeSpecification.exchange(EXCHANGE_NAME))
             .retryWhen(Retry.backoff(NUM_RETRIES, FIRST_BACKOFF));
         Mono<AMQP.Queue.DeclareOk> declareQueue = sender
             .declare(QueueSpecification.queue(QUEUE_NAME)
                 .durable(true)
                 .arguments(rabbitMQConfiguration.workQueueArgumentsBuilder(!AUTO_DELETE_QUEUE)
                     .singleActiveConsumer()
                     .consumerTimeout(rabbitMQConfiguration.getTaskQueueConsumerTimeout().toMillis())
                     .build()))
             .retryWhen(Retry.backoff(NUM_RETRIES, FIRST_BACKOFF));
         Mono<AMQP.Queue.BindOk> bindQueueToExchange = sender
             .bind(BindingSpecification.binding(EXCHANGE_NAME, ROUTING_KEY, QUEUE_NAME))
             .retryWhen(Retry.backoff(NUM_RETRIES, FIRST_BACKOFF));

         declareExchange
             .then(declareQueue)
             .then(bindQueueToExchange)
             .block();
     }

     @Override
     public void restart() {
         closeRabbitResources();
         start();
     }

     private void consumeWorkqueue() {
         receiverHandle = Flux.using(
                 receiverProvider::createReceiver,
                 receiver -> receiver.consumeManualAck(QUEUE_NAME, new ConsumeOptions()),
                 Receiver::close)
             .subscribeOn(Schedulers.elastic())
             .concatMap(this::executeTask)
             .subscribe();
     }

     private Mono<Task.Result> executeTask(AcknowledgableDelivery delivery) {
         return Mono.fromCallable(() -> delivery.getProperties().getHeaders())
             .map(headers -> headers.get(TASK_ID))
             .map(taskIdValue -> TaskId.fromString(taskIdValue.toString()))
             .flatMap(taskId -> Mono.fromCallable(() -> new String(delivery.getBody(), StandardCharsets.UTF_8))
                 .flatMap(bodyValue -> deserialize(bodyValue, taskId))
                 .doOnNext(task -> delivery.ack())
                 .flatMap(task -> executeOnWorker(taskId, task)))
             .onErrorResume(error -> {
                 Optional<Object> taskId = Optional.ofNullable(delivery.getProperties())
                     .flatMap(props -> Optional.ofNullable(props.getHeaders()))
                     .flatMap(headers -> Optional.ofNullable(headers.get(TASK_ID)));
                 LOGGER.error("Unable to process {} {}", TASK_ID, taskId, error);
                 delivery.nack(!REQUEUE);
                 return Mono.empty();
             });
     }

     private Mono<Task> deserialize(String json, TaskId taskId) {
         return Mono.fromCallable(() -> taskSerializer.deserialize(json))
             .onErrorResume(error -> {
                 String errorMessage = String.format("Unable to deserialize submitted Task %s", taskId.asString());
                 LOGGER.error(errorMessage, error);
                 return Mono.from(worker.fail(taskId, Optional.empty(), errorMessage, error))
                     .then(Mono.empty());
             });
     }

     private Mono<Task.Result> executeOnWorker(TaskId taskId, Task task) {
         return worker.executeTask(new TaskWithId(taskId, task))
             .onErrorResume(error -> {
                 String errorMessage = String.format("Unable to run submitted Task %s", taskId.asString());
                 LOGGER.warn(errorMessage, error);
                 return Mono.from(worker.fail(taskId, task.details(), errorMessage, error))
                     .then(Mono.empty());
             });
     }

     private void listenToCancelRequests() {
         sender.declareExchange(ExchangeSpecification.exchange(CANCEL_REQUESTS_EXCHANGE_NAME)).block();
         sender.declare(QueueSpecification.queue(cancelRequestQueueName.asString()).durable(!DURABLE).autoDelete(AUTO_DELETE)).block();
         sender.bind(BindingSpecification.binding(CANCEL_REQUESTS_EXCHANGE_NAME, CANCEL_REQUESTS_ROUTING_KEY, cancelRequestQueueName.asString())).block();
         registerCancelRequestsListener(cancelRequestQueueName.asString());

         sendCancelRequestsQueue = UnicastProcessor.create();
         sendCancelRequestsQueueHandle = sender
             .send(sendCancelRequestsQueue.map(this::makeCancelRequestMessage))
             .subscribeOn(Schedulers.elastic())
             .subscribe();
     }

     private void registerCancelRequestsListener(String queueName) {
         cancelRequestListenerHandle = Flux.using(
                 receiverProvider::createReceiver,
                 receiver -> receiver.consumeAutoAck(queueName),
                 Receiver::close)
             .subscribeOn(Schedulers.elastic())
             .map(this::readCancelRequestMessage)
             .doOnNext(worker::cancelTask)
             .subscribe();
     }

     private TaskId readCancelRequestMessage(Delivery delivery) {
         String message = new String(delivery.getBody(), StandardCharsets.UTF_8);
         return TaskId.fromString(message);
     }

     private OutboundMessage makeCancelRequestMessage(TaskId taskId) {
         byte[] payload = taskId.asString().getBytes(StandardCharsets.UTF_8);
         AMQP.BasicProperties basicProperties = new AMQP.BasicProperties.Builder().build();
         return new OutboundMessage(CANCEL_REQUESTS_EXCHANGE_NAME, CANCEL_REQUESTS_ROUTING_KEY, basicProperties, payload);
     }

     @Override
     public void submit(TaskWithId taskWithId) {
         try {
             byte[] payload = taskSerializer.serialize(taskWithId.getTask()).getBytes(StandardCharsets.UTF_8);
             AMQP.BasicProperties basicProperties = new AMQP.BasicProperties.Builder()
                 .deliveryMode(PERSISTENT_TEXT_PLAIN.getDeliveryMode())
                 .priority(PERSISTENT_TEXT_PLAIN.getPriority())
                 .contentType(PERSISTENT_TEXT_PLAIN.getContentType())
                 .headers(ImmutableMap.of(TASK_ID, taskWithId.getId().asString()))
                 .build();

             OutboundMessage outboundMessage = new OutboundMessage(EXCHANGE_NAME, ROUTING_KEY, basicProperties, payload);
             sender.send(Mono.just(outboundMessage)).block();
         } catch (JsonProcessingException e) {
             throw new RuntimeException(e);
         }
     }

     @Override
     public void cancel(TaskId taskId) {
         sendCancelRequestsQueue.onNext(taskId);
     }

     @Override
     public void close() {
         closeRabbitResources();
     }

     private void closeRabbitResources() {
         Optional.ofNullable(receiverHandle).ifPresent(Disposable::dispose);
         Optional.ofNullable(sendCancelRequestsQueueHandle).ifPresent(Disposable::dispose);
         Optional.ofNullable(cancelRequestListenerHandle).ifPresent(Disposable::dispose);
     }
 }
	/**
	* *************************************************************
	* Licensed to the Apache Software Foundation (ASF) under one *
	* or more contributor license agreements. See the NOTICE file *
	* distributed with this work for additional information *
	* regarding copyright ownership. The ASF licenses this file *
	* to you 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 org.apache.james.task.eventsourcing.distributed;

	import static com.rabbitmq.client.MessageProperties.PERSISTENT_TEXT_PLAIN;
	import static org.apache.james.backends.rabbitmq.Constants.AUTO_DELETE;
	import static org.apache.james.backends.rabbitmq.Constants.DURABLE;
	import static org.apache.james.backends.rabbitmq.Constants.REQUEUE;

	import java.nio.charset.StandardCharsets;
	import java.time.Duration;
	import java.util.Optional;

	import org.apache.james.backends.rabbitmq.RabbitMQConfiguration;
	import org.apache.james.backends.rabbitmq.ReceiverProvider;
	import org.apache.james.server.task.json.JsonTaskSerializer;
	import org.apache.james.task.Task;
	import org.apache.james.task.TaskId;
	import org.apache.james.task.TaskManagerWorker;
	import org.apache.james.task.TaskWithId;
	import org.apache.james.task.WorkQueue;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import com.fasterxml.jackson.core.JsonProcessingException;
	import com.google.common.annotations.VisibleForTesting;
	import com.google.common.collect.ImmutableMap;
	import com.rabbitmq.client.AMQP;
	import com.rabbitmq.client.Delivery;

	import reactor.core.Disposable;
	import reactor.core.publisher.Flux;
	import reactor.core.publisher.Mono;
	import reactor.core.publisher.UnicastProcessor;
	import reactor.core.scheduler.Schedulers;
	import reactor.rabbitmq.AcknowledgableDelivery;
	import reactor.rabbitmq.BindingSpecification;
	import reactor.rabbitmq.ConsumeOptions;
	import reactor.rabbitmq.ExchangeSpecification;
	import reactor.rabbitmq.OutboundMessage;
	import reactor.rabbitmq.QueueSpecification;
	import reactor.rabbitmq.Receiver;
	import reactor.rabbitmq.Sender;
	import reactor.util.retry.Retry;

	public class RabbitMQWorkQueue implements WorkQueue {
	private static final Logger LOGGER = LoggerFactory.getLogger(RabbitMQWorkQueue.class);

	static final String EXCHANGE_NAME = "taskManagerWorkQueueExchange";
	static final String QUEUE_NAME = "taskManagerWorkQueue";
	static final String ROUTING_KEY = "taskManagerWorkQueueRoutingKey";

	static final String CANCEL_REQUESTS_EXCHANGE_NAME = "taskManagerCancelRequestsExchange";
	static final String CANCEL_REQUESTS_ROUTING_KEY = "taskManagerCancelRequestsRoutingKey";
	public static final String TASK_ID = "taskId";

	public static final int NUM_RETRIES = 8;
	public static final Duration FIRST_BACKOFF = Duration.ofMillis(100);
	public static final boolean AUTO_DELETE_QUEUE = true;

	private final TaskManagerWorker worker;
	private final JsonTaskSerializer taskSerializer;
	private final RabbitMQWorkQueueConfiguration configuration;
	private final Sender sender;
	private final ReceiverProvider receiverProvider;
	private final RabbitMQConfiguration rabbitMQConfiguration;
	private final CancelRequestQueueName cancelRequestQueueName;
	private UnicastProcessor<TaskId> sendCancelRequestsQueue;
	private Disposable sendCancelRequestsQueueHandle;
	private Disposable receiverHandle;
	private Disposable cancelRequestListenerHandle;

	public RabbitMQWorkQueue(TaskManagerWorker worker, Sender sender,
	ReceiverProvider receiverProvider, JsonTaskSerializer taskSerializer,
	RabbitMQWorkQueueConfiguration configuration, CancelRequestQueueName cancelRequestQueueName,
	RabbitMQConfiguration rabbitMQConfiguration) {
	this.cancelRequestQueueName = cancelRequestQueueName;
	this.worker = worker;
	this.receiverProvider = receiverProvider;
	this.sender = sender;
	this.taskSerializer = taskSerializer;
	this.configuration = configuration;
	this.rabbitMQConfiguration = rabbitMQConfiguration;
	}

	@Override
	public void start() {
	startWorkqueue();
	listenToCancelRequests();
	}

	private void startWorkqueue() {
	declareQueue();

	if (configuration.enabled()) {
	consumeWorkqueue();
	}
	}

	@VisibleForTesting
	void declareQueue() {
	Mono<AMQP.Exchange.DeclareOk> declareExchange = sender
	.declareExchange(ExchangeSpecification.exchange(EXCHANGE_NAME))
	.retryWhen(Retry.backoff(NUM_RETRIES, FIRST_BACKOFF));
	Mono<AMQP.Queue.DeclareOk> declareQueue = sender
	.declare(QueueSpecification.queue(QUEUE_NAME)
	.durable(true)
	.arguments(rabbitMQConfiguration.workQueueArgumentsBuilder(!AUTO_DELETE_QUEUE)
	.singleActiveConsumer()
	.consumerTimeout(rabbitMQConfiguration.getTaskQueueConsumerTimeout().toMillis())
	.build()))
	.retryWhen(Retry.backoff(NUM_RETRIES, FIRST_BACKOFF));
	Mono<AMQP.Queue.BindOk> bindQueueToExchange = sender
	.bind(BindingSpecification.binding(EXCHANGE_NAME, ROUTING_KEY, QUEUE_NAME))
	.retryWhen(Retry.backoff(NUM_RETRIES, FIRST_BACKOFF));

	declareExchange
	.then(declareQueue)
	.then(bindQueueToExchange)
	.block();
	}

	@Override
	public void restart() {
	closeRabbitResources();
	start();
	}

	private void consumeWorkqueue() {
	receiverHandle = Flux.using(
	receiverProvider::createReceiver,
	receiver -> receiver.consumeManualAck(QUEUE_NAME, new ConsumeOptions()),
	Receiver::close)
	.subscribeOn(Schedulers.elastic())
	.concatMap(this::executeTask)
	.subscribe();
	}

	private Mono<Task.Result> executeTask(AcknowledgableDelivery delivery) {
	return Mono.fromCallable(() -> delivery.getProperties().getHeaders())
	.map(headers -> headers.get(TASK_ID))
	.map(taskIdValue -> TaskId.fromString(taskIdValue.toString()))
	.flatMap(taskId -> Mono.fromCallable(() -> new String(delivery.getBody(), StandardCharsets.UTF_8))
	.flatMap(bodyValue -> deserialize(bodyValue, taskId))
	.doOnNext(task -> delivery.ack())
	.flatMap(task -> executeOnWorker(taskId, task)))
	.onErrorResume(error -> {
	Optional<Object> taskId = Optional.ofNullable(delivery.getProperties())
	.flatMap(props -> Optional.ofNullable(props.getHeaders()))
	.flatMap(headers -> Optional.ofNullable(headers.get(TASK_ID)));
	LOGGER.error("Unable to process {} {}", TASK_ID, taskId, error);
	delivery.nack(!REQUEUE);
	return Mono.empty();
	});
	}

	private Mono<Task> deserialize(String json, TaskId taskId) {
	return Mono.fromCallable(() -> taskSerializer.deserialize(json))
	.onErrorResume(error -> {
	String errorMessage = String.format("Unable to deserialize submitted Task %s", taskId.asString());
	LOGGER.error(errorMessage, error);
	return Mono.from(worker.fail(taskId, Optional.empty(), errorMessage, error))
	.then(Mono.empty());
	});
	}

	private Mono<Task.Result> executeOnWorker(TaskId taskId, Task task) {
	return worker.executeTask(new TaskWithId(taskId, task))
	.onErrorResume(error -> {
	String errorMessage = String.format("Unable to run submitted Task %s", taskId.asString());
	LOGGER.warn(errorMessage, error);
	return Mono.from(worker.fail(taskId, task.details(), errorMessage, error))
	.then(Mono.empty());
	});
	}

	private void listenToCancelRequests() {
	sender.declareExchange(ExchangeSpecification.exchange(CANCEL_REQUESTS_EXCHANGE_NAME)).block();
	sender.declare(QueueSpecification.queue(cancelRequestQueueName.asString()).durable(!DURABLE).autoDelete(AUTO_DELETE)).block();
	sender.bind(BindingSpecification.binding(CANCEL_REQUESTS_EXCHANGE_NAME, CANCEL_REQUESTS_ROUTING_KEY, cancelRequestQueueName.asString())).block();
	registerCancelRequestsListener(cancelRequestQueueName.asString());

	sendCancelRequestsQueue = UnicastProcessor.create();
	sendCancelRequestsQueueHandle = sender
	.send(sendCancelRequestsQueue.map(this::makeCancelRequestMessage))
	.subscribeOn(Schedulers.elastic())
	.subscribe();
	}

	private void registerCancelRequestsListener(String queueName) {
	cancelRequestListenerHandle = Flux.using(
	receiverProvider::createReceiver,
	receiver -> receiver.consumeAutoAck(queueName),
	Receiver::close)
	.subscribeOn(Schedulers.elastic())
	.map(this::readCancelRequestMessage)
	.doOnNext(worker::cancelTask)
	.subscribe();
	}

	private TaskId readCancelRequestMessage(Delivery delivery) {
	String message = new String(delivery.getBody(), StandardCharsets.UTF_8);
	return TaskId.fromString(message);
	}

	private OutboundMessage makeCancelRequestMessage(TaskId taskId) {
	byte[] payload = taskId.asString().getBytes(StandardCharsets.UTF_8);
	AMQP.BasicProperties basicProperties = new AMQP.BasicProperties.Builder().build();
	return new OutboundMessage(CANCEL_REQUESTS_EXCHANGE_NAME, CANCEL_REQUESTS_ROUTING_KEY, basicProperties, payload);
	}

	@Override
	public void submit(TaskWithId taskWithId) {
	try {
	byte[] payload = taskSerializer.serialize(taskWithId.getTask()).getBytes(StandardCharsets.UTF_8);
	AMQP.BasicProperties basicProperties = new AMQP.BasicProperties.Builder()
	.deliveryMode(PERSISTENT_TEXT_PLAIN.getDeliveryMode())
	.priority(PERSISTENT_TEXT_PLAIN.getPriority())
	.contentType(PERSISTENT_TEXT_PLAIN.getContentType())
	.headers(ImmutableMap.of(TASK_ID, taskWithId.getId().asString()))
	.build();

	OutboundMessage outboundMessage = new OutboundMessage(EXCHANGE_NAME, ROUTING_KEY, basicProperties, payload);
	sender.send(Mono.just(outboundMessage)).block();
	} catch (JsonProcessingException e) {
	throw new RuntimeException(e);
	}
	}

	@Override
	public void cancel(TaskId taskId) {
	sendCancelRequestsQueue.onNext(taskId);
	}

	@Override
	public void close() {
	closeRabbitResources();
	}

	private void closeRabbitResources() {
	Optional.ofNullable(receiverHandle).ifPresent(Disposable::dispose);
	Optional.ofNullable(sendCancelRequestsQueueHandle).ifPresent(Disposable::dispose);
	Optional.ofNullable(cancelRequestListenerHandle).ifPresent(Disposable::dispose);
	}
	}