org.apache.helix.ipc.HelixIPCMessageManager Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of helix-ipc Show documentation
The newest version!
package org.apache.helix.ipc;

/*
 * 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.
 */

import io.netty.buffer.ByteBuf;
import org.apache.helix.resolver.HelixAddress;
import org.apache.helix.resolver.HelixMessageScope;
import org.apache.log4j.Logger;

import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * A wrapper around a base IPC service that manages message retries / timeouts.
 * 
 * This class manages retries and timeouts, and can be used in the same way as a
 * {@link org.apache.helix.ipc.HelixIPCService}.
 * 
 * 
 * A message will be sent until the max number of retries has been reached (i.e. timed out), or it
 * is acknowledged by the recipient. If the max number of retries is -1, it will be retried forever.
 * 
 * 
 * A callback should be registered for every acknowledgement message type associated with any
 * original message type sent by this class.
 * 
 * 
 * For example, consider we have the two message types defined: DATA_REQ = 1, DATA_ACK = 2. One
 * would do the following:
 * 
 * 
 * messageManager.registerCallback(DATA_ACK, new HelixIPCCallback() {
 *   public void onMessage(HelixMessageScope scope, UUID messageId, ByteBuf message) {
 *     // Process ACK
 *   }
 * 
 *   public void onError(HelixMessageScope scope, UUID messageId, Throwable cause) {
 *     // Message error or timeout
 *   }
 * });
 * 
 * messageManager.send(destinations, DATA_REQ, messageId, data);
 * 
 * 
 * 
 * 
 * In send, we note messageId, and retry until we get a DATA_ACK for the same messageId. The
 * callback registered with the message manager will only be called once, even if the message is
 * acknowledged several times.
 * 
 */
public class HelixIPCMessageManager implements HelixIPCService {

  private static final Logger LOG = Logger.getLogger(HelixIPCMessageManager.class);

  private final ScheduledExecutorService scheduler;
  private final HelixIPCService baseIpcService;
  private final long messageTimeoutMillis;
  private final int maxNumRetries;
  private final ConcurrentMap pendingMessages;
  private final ConcurrentMap retriesLeft;
  private final ConcurrentMap messageBuffers;
  private final ConcurrentMap callbacks;

  public HelixIPCMessageManager(ScheduledExecutorService scheduler, HelixIPCService baseIpcService,
      long messageTimeoutMillis, int maxNumRetries) {
    this.scheduler = scheduler;
    this.baseIpcService = baseIpcService;
    this.maxNumRetries = maxNumRetries;
    this.messageTimeoutMillis = messageTimeoutMillis;
    this.pendingMessages = new ConcurrentHashMap();
    this.retriesLeft = new ConcurrentHashMap();
    this.messageBuffers = new ConcurrentHashMap();
    this.callbacks = new ConcurrentHashMap();
  }

  @Override
  public void start() throws Exception {
    baseIpcService.start();
  }

  @Override
  public void shutdown() throws Exception {
    baseIpcService.shutdown();
  }

  @Override
  public void send(final HelixAddress destination, final int messageType, final UUID messageId,
      final ByteBuf message) {
    // State
    pendingMessages.put(messageId, true);
    retriesLeft.putIfAbsent(messageId, new AtomicInteger(maxNumRetries));
    messageBuffers.put(messageId, message);

    // Will free it when we've finally received response
    message.retain();

    // Send initial message
    baseIpcService.send(destination, messageType, messageId, message);

    // Retries
    scheduler.schedule(new Runnable() {
      @Override
      public void run() {
        Boolean isPending = pendingMessages.get(messageId);
        AtomicInteger numLeft = retriesLeft.get(messageId);
        if (numLeft != null && isPending != null && isPending) {
          if (numLeft.decrementAndGet() > 0 || maxNumRetries == -1) {
            // n.b. will schedule another retry
            send(destination, messageType, messageId, message);
          } else {
            LOG.warn("Message " + messageId + " timed out after " + maxNumRetries + " retries");
          }
        }
      }
    }, messageTimeoutMillis, TimeUnit.MILLISECONDS);
  }

  @Override
  public void registerCallback(final int messageType, final HelixIPCCallback callback) {

    // This callback will first check if the message is pending, then delegate to the provided
    // callback if it has not yet done so.
    HelixIPCCallback wrappedCallback = new HelixIPCCallback() {
      @Override
      public void onMessage(HelixMessageScope scope, UUID messageId, ByteBuf message) {
        if (pendingMessages.replace(messageId, true, false)) {
          pendingMessages.remove(messageId);
          ByteBuf originalMessage = messageBuffers.remove(messageId);
          if (originalMessage != null) {
            originalMessage.release();
          }
          retriesLeft.remove(messageId);
          callback.onMessage(scope, messageId, message);
        }
      }
    };

    callbacks.put(messageType, wrappedCallback);
    baseIpcService.registerCallback(messageType, wrappedCallback);
  }
}