io.atomix.cluster.messaging.impl.DefaultClusterCommunicationService Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of zeebe-atomix-cluster Show documentation
The newest version!
/*
 * Copyright 2017-present Open Networking Foundation
 * Copyright © 2020 camunda services GmbH ([email protected])
 *
 * 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 io.atomix.cluster.messaging.impl;

import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.base.Objects;
import com.google.common.collect.Maps;
import io.atomix.cluster.ClusterMembershipService;
import io.atomix.cluster.Member;
import io.atomix.cluster.MemberId;
import io.atomix.cluster.messaging.ClusterCommunicationService;
import io.atomix.cluster.messaging.ManagedClusterCommunicationService;
import io.atomix.cluster.messaging.MessagingException;
import io.atomix.cluster.messaging.MessagingException.NoSuchMemberException;
import io.atomix.cluster.messaging.MessagingService;
import io.atomix.cluster.messaging.UnicastService;
import io.atomix.utils.net.Address;
import java.time.Duration;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Cluster communication service implementation. */
public class DefaultClusterCommunicationService implements ManagedClusterCommunicationService {

  protected final ClusterMembershipService membershipService;
  protected final MessagingService messagingService;
  protected final UnicastService unicastService;
  private final Logger log = LoggerFactory.getLogger(getClass());
  private final Map> unicastConsumers = Maps.newConcurrentMap();
  private final AtomicBoolean started = new AtomicBoolean();

  public DefaultClusterCommunicationService(
      final ClusterMembershipService membershipService,
      final MessagingService messagingService,
      final UnicastService unicastService) {
    this.membershipService = checkNotNull(membershipService, "clusterService cannot be null");
    this.messagingService = checkNotNull(messagingService, "messagingService cannot be null");
    this.unicastService = checkNotNull(unicastService, "unicastService cannot be null");
  }

  @Override
  public  void broadcast(
      final String subject,
      final M message,
      final Function encoder,
      final boolean reliable) {
    multicast(
        subject,
        message,
        encoder,
        membershipService.getMembers().stream()
            .filter(node -> !Objects.equal(node, membershipService.getLocalMember()))
            .map(Member::id)
            .collect(Collectors.toSet()),
        reliable);
  }

  @Override
  public  void multicast(
      final String subject,
      final M message,
      final Function encoder,
      final Set nodes,
      final boolean reliable) {
    final byte[] payload = encoder.apply(message);
    nodes.forEach(memberId -> doUnicast(subject, payload, memberId, reliable));
  }

  @Override
  public  void unicast(
      final String subject,
      final M message,
      final Function encoder,
      final MemberId memberId,
      final boolean reliable) {
    final byte[] payload = encoder.apply(message);
    doUnicast(subject, payload, memberId, reliable);
  }

  @Override
  public  CompletableFuture send(
      final String subject,
      final M message,
      final Function encoder,
      final Function decoder,
      final MemberId toMemberId,
      final Duration timeout) {
    try {
      return sendAndReceive(subject, encoder.apply(message), toMemberId, timeout)
          .thenApply(decoder);
    } catch (final Exception e) {
      return CompletableFuture.failedFuture(e);
    }
  }

  @Override
  public  void replyTo(
      final String subject,
      final Function decoder,
      final Function> handler,
      final Function encoder) {
    messagingService.registerHandler(
        subject, new InternalMessageResponder<>(decoder, encoder, handler));
  }

  @Override
  public  void consume(
      final String subject,
      final Function decoder,
      final Consumer handler,
      final Executor executor) {
    messagingService.registerHandler(
        subject, new InternalMessageConsumer<>(decoder, handler), executor);
    final BiConsumer unicastConsumer =
        new InternalMessageConsumer<>(decoder, handler);
    unicastConsumers.put(subject, unicastConsumer);
    unicastService.addListener(subject, unicastConsumer, executor);
  }

  @Override
  public  void consume(
      final String subject,
      final Function decoder,
      final BiConsumer handler,
      final Executor executor) {
    messagingService.registerHandler(
        subject, new InternalMessageBiConsumer<>(decoder, handler), executor);
    final BiConsumer unicastConsumer =
        new InternalMessageBiConsumer<>(decoder, handler);
    unicastConsumers.put(subject, unicastConsumer);
    unicastService.addListener(subject, unicastConsumer, executor);
  }

  @Override
  public  void replyTo(
      final String subject,
      final Function decoder,
      final BiFunction handler,
      final Function encoder,
      final Executor executor) {
    messagingService.registerHandler(
        subject, new InternalMessageBiResponder<>(decoder, encoder, handler, executor));
  }

  @Override
  public  void replyToAsync(
      final String subject,
      final Function decoder,
      final Function> handler,
      final Function encoder,
      final Executor executor) {
    messagingService.registerHandler(
        subject, new InternalMessageAsyncResponder<>(decoder, encoder, handler, executor));
  }

  @Override
  public void unsubscribe(final String subject) {
    messagingService.unregisterHandler(subject);
    final BiConsumer consumer = unicastConsumers.remove(subject);
    if (consumer != null) {
      unicastService.removeListener(subject, consumer);
    }
  }

  private void doUnicast(
      final String subject,
      final byte[] payload,
      final MemberId toMemberId,
      final boolean reliable) {
    final Member member = membershipService.getMember(toMemberId);
    if (member == null) {
      return;
    }

    if (reliable) {
      messagingService.sendAsync(member.address(), subject, payload);
    } else {
      unicastService.unicast(member.address(), subject, payload);
    }
  }

  private CompletableFuture sendAndReceive(
      final String subject,
      final byte[] payload,
      final MemberId toMemberId,
      final Duration timeout) {
    final Member member = membershipService.getMember(toMemberId);
    if (member == null) {
      return failOnMemberNotKnown(subject, toMemberId);
    }

    return messagingService.sendAndReceive(member.address(), subject, payload, timeout);
  }

  private  CompletableFuture failOnMemberNotKnown(
      final String subject, final MemberId toMemberId) {
    final var errorMessage =
        String.format(
            "Expected to send a message with subject '%s' to member '%s', but member is not known. Known members are '%s'.",
            subject, toMemberId, membershipService.getMembers());
    return CompletableFuture.failedFuture(new NoSuchMemberException(errorMessage));
  }

  @Override
  public CompletableFuture start() {
    if (started.compareAndSet(false, true)) {
      log.info("Started");
    }
    return CompletableFuture.completedFuture(this);
  }

  @Override
  public boolean isRunning() {
    return started.get();
  }

  @Override
  public CompletableFuture stop() {
    if (started.compareAndSet(true, false)) {
      log.info("Stopped");
    }
    return CompletableFuture.completedFuture(null);
  }

  private static class InternalMessageConsumer implements BiConsumer {

    private final Function decoder;
    private final Consumer consumer;

    InternalMessageConsumer(final Function decoder, final Consumer consumer) {
      this.decoder = decoder;
      this.consumer = consumer;
    }

    @Override
    public void accept(final Address sender, final byte[] bytes) {
      consumer.accept(decoder.apply(bytes));
    }
  }

  private static class InternalMessageResponder
      implements BiFunction> {
    private final Function decoder;
    private final Function encoder;
    private final Function> handler;

    InternalMessageResponder(
        final Function decoder,
        final Function encoder,
        final Function> handler) {
      this.decoder = decoder;
      this.encoder = encoder;
      this.handler = handler;
    }

    @Override
    public CompletableFuture apply(final Address sender, final byte[] bytes) {
      return handler.apply(decoder.apply(bytes)).thenApply(encoder);
    }
  }

  private static class InternalMessageAsyncResponder
      implements BiFunction> {
    private final Function decoder;
    private final Function encoder;
    private final Function> handler;
    private final Executor executor;

    InternalMessageAsyncResponder(
        final Function decoder,
        final Function encoder,
        final Function> handler,
        final Executor executor) {
      this.decoder = decoder;
      this.encoder = encoder;
      this.handler = handler;
      this.executor = executor;
    }

    @Override
    public CompletableFuture apply(final Address sender, final byte[] bytes) {
      return CompletableFuture.supplyAsync(() -> decoder.apply(bytes), executor)
          .thenComposeAsync(handler, executor)
          .thenApplyAsync(encoder, executor);
    }
  }

  private class InternalMessageBiConsumer implements BiConsumer {

    private final Function decoder;
    private final BiConsumer consumer;

    InternalMessageBiConsumer(
        final Function decoder, final BiConsumer consumer) {
      this.decoder = decoder;
      this.consumer = consumer;
    }

    @Override
    public void accept(final Address sender, final byte[] bytes) {
      final Member member = membershipService.getMember(sender);
      if (member != null) {
        consumer.accept(member.id(), decoder.apply(bytes));
      }
    }
  }

  private final class InternalMessageBiResponder
      implements BiFunction> {
    private final Function decoder;
    private final Function encoder;
    private final BiFunction handler;
    private final Executor executor;

    InternalMessageBiResponder(
        final Function decoder,
        final Function encoder,
        final BiFunction handler,
        final Executor executor) {
      this.decoder = decoder;
      this.encoder = encoder;
      this.handler = handler;
      this.executor = executor;
    }

    @Override
    public CompletableFuture apply(final Address address, final byte[] bytes) {
      final var response = new CompletableFuture();
      executor.execute(
          () -> {
            try {
              handleRequest(address, bytes, response);
            } catch (final Exception e) {
              response.completeExceptionally(e);
            }
          });

      return response;
    }

    private void handleRequest(
        final Address address, final byte[] bytes, final CompletableFuture response)
        throws Exception {
      final Member member = membershipService.getMember(address);
      if (member == null) {
        throw new MessagingException.NoSuchMemberException(address);
      }

      final var decodedResponse = handler.apply(member.id(), decoder.apply(bytes));
      final var encodedResponse = encoder.apply(decodedResponse);
      response.complete(encodedResponse);
    }
  }
}