helix-core/src/main/java/org/apache/helix/messaging/handling/BatchMessageHandler.java - helix - Git at Google

 package org.apache.helix.messaging.handling;

 /*
  * 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 java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;

 import org.apache.helix.HelixDataAccessor;
 import org.apache.helix.HelixException;
 import org.apache.helix.HelixManager;
 import org.apache.helix.NotificationContext;
 import org.apache.helix.NotificationContext.MapKey;
 import org.apache.helix.PropertyKey;
 import org.apache.helix.model.CurrentState;
 import org.apache.helix.model.Message;
 import org.apache.helix.model.Message.Attributes;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 public class BatchMessageHandler extends MessageHandler {
   private static Logger LOG = LoggerFactory.getLogger(BatchMessageHandler.class);

   final MessageHandlerFactory _msgHandlerFty;
   final TaskExecutor _executor;
   final List<Message> _subMessages;
   final List<MessageHandler> _subMessageHandlers;
   final BatchMessageWrapper _batchMsgWrapper;

   public BatchMessageHandler(Message msg, NotificationContext context, MessageHandlerFactory fty,
       BatchMessageWrapper wrapper, TaskExecutor executor) {
     super(msg, context);

     if (fty == null || executor == null) {
       throw new HelixException("MessageHandlerFactory | TaskExecutor can't be null");
     }

     _msgHandlerFty = fty;
     _batchMsgWrapper = wrapper;
     _executor = executor;

     // create sub-messages
     _subMessages = new ArrayList<Message>();
     List<String> partitionKeys = _message.getPartitionNames();
     for (String partitionKey : partitionKeys) {
       // assign a new message id, put batch-msg-id to parent-id field
       Message subMsg = new Message(_message.getRecord(), UUID.randomUUID().toString());
       subMsg.setPartitionName(partitionKey);
       subMsg.setAttribute(Attributes.PARENT_MSG_ID, _message.getId());
       subMsg.setBatchMessageMode(false);

       _subMessages.add(subMsg);
     }

     // create sub-message handlers
     _subMessageHandlers = createMsgHandlers(_subMessages, context);
   }

   List<MessageHandler> createMsgHandlers(List<Message> msgs, NotificationContext context) {

     List<MessageHandler> handlers = new ArrayList<MessageHandler>();
     for (Message msg : msgs) {
       MessageHandler handler = _msgHandlerFty.createHandler(msg, context);
       handlers.add(handler);
     }
     return handlers;
   }

   public void preHandleMessage() {
     if (_message.getBatchMessageMode() == true && _batchMsgWrapper != null) {
       _batchMsgWrapper.start(_message, _notificationContext);
     }
   }

   public void postHandleMessage() {
     if (_message.getBatchMessageMode() == true && _batchMsgWrapper != null) {
       _batchMsgWrapper.end(_message, _notificationContext);
     }

     // update currentState
     HelixManager manager = _notificationContext.getManager();
     HelixDataAccessor accessor = manager.getHelixDataAccessor();
     ConcurrentHashMap<String, CurrentStateUpdate> csUpdateMap =
         (ConcurrentHashMap<String, CurrentStateUpdate>) _notificationContext
             .get(MapKey.CURRENT_STATE_UPDATE.toString());

     if (csUpdateMap != null) {
       Map<PropertyKey, CurrentState> csUpdate = mergeCurStateUpdate(csUpdateMap);

       // TODO: change to use asyncSet
       for (PropertyKey key : csUpdate.keySet()) {
         // logger.info("updateCS: " + key);
         // System.out.println("\tupdateCS: " + key.getPath() + ", " +
         // curStateMap.get(key));
         if(!accessor.updateProperty(key, csUpdate.get(key))) {
           LOG.error(
               "Fails to persist current state to ZK for key " + key);
         }
       }
     }
   }

   // will not return until all sub-message executions are done
   @Override
   public HelixTaskResult handleMessage() {
     HelixTaskResult result = null;
     List<Future<HelixTaskResult>> futures = null;
     List<MessageTask> batchTasks = new ArrayList<MessageTask>();

     synchronized (_batchMsgWrapper) {
       try {
         preHandleMessage();

         int exeBatchSize = 1; // TODO: getExeBatchSize from msg
         List<String> partitionKeys = _message.getPartitionNames();
         for (int i = 0; i < partitionKeys.size(); i += exeBatchSize) {
           if (i + exeBatchSize <= partitionKeys.size()) {
             List<Message> msgs = _subMessages.subList(i, i + exeBatchSize);
             List<MessageHandler> handlers = _subMessageHandlers.subList(i, i + exeBatchSize);
             HelixBatchMessageTask batchTask =
                 new HelixBatchMessageTask(_message, msgs, handlers, _notificationContext);
             batchTasks.add(batchTask);

           } else {
             List<Message> msgs = _subMessages.subList(i, i + partitionKeys.size());
             List<MessageHandler> handlers =
                 _subMessageHandlers.subList(i, i + partitionKeys.size());

             HelixBatchMessageTask batchTask =
                 new HelixBatchMessageTask(_message, msgs, handlers, _notificationContext);
             batchTasks.add(batchTask);
           }
         }

         // invokeAll() is blocking call
         long timeout = _message.getExecutionTimeout();
         if (timeout == -1) {
           timeout = Long.MAX_VALUE;
         }
         futures = _executor.invokeAllTasks(batchTasks, timeout, TimeUnit.MILLISECONDS);
       } catch (Exception e) {
         LOG.error("fail to execute batchMsg: " + _message.getId(), e);
         result = new HelixTaskResult();
         result.setException(e);

         // HelixTask will call onError on this batch-msg-handler
         // return result;
       }

       // combine sub-results to result
       if (futures != null) {
         boolean isBatchTaskSucceed = true;

         for (int i = 0; i < futures.size(); i++) {
           Future<HelixTaskResult> future = futures.get(i);
           MessageTask subTask = batchTasks.get(i);
           try {
             HelixTaskResult subTaskResult = future.get();
             if (!subTaskResult.isSuccess()) {
               isBatchTaskSucceed = false;
             }
           } catch (InterruptedException e) {
             isBatchTaskSucceed = false;
             LOG.error("interrupted in executing batch-msg: " + _message.getId() + ", sub-msg: "
                 + subTask.getTaskId(), e);
           } catch (ExecutionException e) {
             isBatchTaskSucceed = false;
             LOG.error(
                 "fail to execute batch-msg: " + _message.getId() + ", sub-msg: "
                     + subTask.getTaskId(), e);
           }
         }
         result = new HelixTaskResult();
         result.setSuccess(isBatchTaskSucceed);
       }

       result.setCompleteTime(System.currentTimeMillis());
       // pass task-result to post-handle-msg
       _notificationContext.add(MapKey.HELIX_TASK_RESULT.toString(), result);
       postHandleMessage();

       return result;
     }
   }

   @Override
   public void onError(Exception e, ErrorCode code, ErrorType type) {
     // if one sub-message execution fails, call onError on all sub-message handlers
     for (MessageHandler handler : _subMessageHandlers) {
       handler.onError(e, code, type);
     }
   }

   // TODO: optimize this based on the fact that each cs update is for a
   // distinct partition
   private Map<PropertyKey, CurrentState> mergeCurStateUpdate(
       ConcurrentHashMap<String, CurrentStateUpdate> csUpdateMap) {
     Map<String, CurrentStateUpdate> curStateUpdateMap = new HashMap<String, CurrentStateUpdate>();
     for (CurrentStateUpdate update : csUpdateMap.values()) {
       String path = update._key.getPath(); // TODO: this is time
                                            // consuming, optimize it
       if (!curStateUpdateMap.containsKey(path)) {
         curStateUpdateMap.put(path, update);
       } else {
         // long start = System.currentTimeMillis();
         curStateUpdateMap.get(path).merge(update._delta);
         // long end = System.currentTimeMillis();
         // LOG.info("each merge took: " + (end - start));
       }
     }

     Map<PropertyKey, CurrentState> ret = new HashMap<PropertyKey, CurrentState>();
     for (CurrentStateUpdate update : curStateUpdateMap.values()) {
       ret.put(update._key, update._delta);
     }

     return ret;
   }

 }
	package org.apache.helix.messaging.handling;

	/*
	* 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 java.util.ArrayList;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.UUID;
	import java.util.concurrent.ConcurrentHashMap;
	import java.util.concurrent.ExecutionException;
	import java.util.concurrent.Future;
	import java.util.concurrent.TimeUnit;

	import org.apache.helix.HelixDataAccessor;
	import org.apache.helix.HelixException;
	import org.apache.helix.HelixManager;
	import org.apache.helix.NotificationContext;
	import org.apache.helix.NotificationContext.MapKey;
	import org.apache.helix.PropertyKey;
	import org.apache.helix.model.CurrentState;
	import org.apache.helix.model.Message;
	import org.apache.helix.model.Message.Attributes;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	public class BatchMessageHandler extends MessageHandler {
	private static Logger LOG = LoggerFactory.getLogger(BatchMessageHandler.class);

	final MessageHandlerFactory _msgHandlerFty;
	final TaskExecutor _executor;
	final List<Message> _subMessages;
	final List<MessageHandler> _subMessageHandlers;
	final BatchMessageWrapper _batchMsgWrapper;

	public BatchMessageHandler(Message msg, NotificationContext context, MessageHandlerFactory fty,
	BatchMessageWrapper wrapper, TaskExecutor executor) {
	super(msg, context);

	if (fty == null \|\| executor == null) {
	throw new HelixException("MessageHandlerFactory \| TaskExecutor can't be null");
	}

	_msgHandlerFty = fty;
	_batchMsgWrapper = wrapper;
	_executor = executor;

	// create sub-messages
	_subMessages = new ArrayList<Message>();
	List<String> partitionKeys = _message.getPartitionNames();
	for (String partitionKey : partitionKeys) {
	// assign a new message id, put batch-msg-id to parent-id field
	Message subMsg = new Message(_message.getRecord(), UUID.randomUUID().toString());
	subMsg.setPartitionName(partitionKey);
	subMsg.setAttribute(Attributes.PARENT_MSG_ID, _message.getId());
	subMsg.setBatchMessageMode(false);

	_subMessages.add(subMsg);
	}

	// create sub-message handlers
	_subMessageHandlers = createMsgHandlers(_subMessages, context);
	}

	List<MessageHandler> createMsgHandlers(List<Message> msgs, NotificationContext context) {

	List<MessageHandler> handlers = new ArrayList<MessageHandler>();
	for (Message msg : msgs) {
	MessageHandler handler = _msgHandlerFty.createHandler(msg, context);
	handlers.add(handler);
	}
	return handlers;
	}

	public void preHandleMessage() {
	if (_message.getBatchMessageMode() == true && _batchMsgWrapper != null) {
	_batchMsgWrapper.start(_message, _notificationContext);
	}
	}

	public void postHandleMessage() {
	if (_message.getBatchMessageMode() == true && _batchMsgWrapper != null) {
	_batchMsgWrapper.end(_message, _notificationContext);
	}

	// update currentState
	HelixManager manager = _notificationContext.getManager();
	HelixDataAccessor accessor = manager.getHelixDataAccessor();
	ConcurrentHashMap<String, CurrentStateUpdate> csUpdateMap =
	(ConcurrentHashMap<String, CurrentStateUpdate>) _notificationContext
	.get(MapKey.CURRENT_STATE_UPDATE.toString());

	if (csUpdateMap != null) {
	Map<PropertyKey, CurrentState> csUpdate = mergeCurStateUpdate(csUpdateMap);

	// TODO: change to use asyncSet
	for (PropertyKey key : csUpdate.keySet()) {
	// logger.info("updateCS: " + key);
	// System.out.println("\tupdateCS: " + key.getPath() + ", " +
	// curStateMap.get(key));
	if(!accessor.updateProperty(key, csUpdate.get(key))) {
	LOG.error(
	"Fails to persist current state to ZK for key " + key);
	}
	}
	}
	}

	// will not return until all sub-message executions are done
	@Override
	public HelixTaskResult handleMessage() {
	HelixTaskResult result = null;
	List<Future<HelixTaskResult>> futures = null;
	List<MessageTask> batchTasks = new ArrayList<MessageTask>();

	synchronized (_batchMsgWrapper) {
	try {
	preHandleMessage();

	int exeBatchSize = 1; // TODO: getExeBatchSize from msg
	List<String> partitionKeys = _message.getPartitionNames();
	for (int i = 0; i < partitionKeys.size(); i += exeBatchSize) {
	if (i + exeBatchSize <= partitionKeys.size()) {
	List<Message> msgs = _subMessages.subList(i, i + exeBatchSize);
	List<MessageHandler> handlers = _subMessageHandlers.subList(i, i + exeBatchSize);
	HelixBatchMessageTask batchTask =
	new HelixBatchMessageTask(_message, msgs, handlers, _notificationContext);
	batchTasks.add(batchTask);

	} else {
	List<Message> msgs = _subMessages.subList(i, i + partitionKeys.size());
	List<MessageHandler> handlers =
	_subMessageHandlers.subList(i, i + partitionKeys.size());

	HelixBatchMessageTask batchTask =
	new HelixBatchMessageTask(_message, msgs, handlers, _notificationContext);
	batchTasks.add(batchTask);
	}
	}

	// invokeAll() is blocking call
	long timeout = _message.getExecutionTimeout();
	if (timeout == -1) {
	timeout = Long.MAX_VALUE;
	}
	futures = _executor.invokeAllTasks(batchTasks, timeout, TimeUnit.MILLISECONDS);
	} catch (Exception e) {
	LOG.error("fail to execute batchMsg: " + _message.getId(), e);
	result = new HelixTaskResult();
	result.setException(e);

	// HelixTask will call onError on this batch-msg-handler
	// return result;
	}

	// combine sub-results to result
	if (futures != null) {
	boolean isBatchTaskSucceed = true;

	for (int i = 0; i < futures.size(); i++) {
	Future<HelixTaskResult> future = futures.get(i);
	MessageTask subTask = batchTasks.get(i);
	try {
	HelixTaskResult subTaskResult = future.get();
	if (!subTaskResult.isSuccess()) {
	isBatchTaskSucceed = false;
	}
	} catch (InterruptedException e) {
	isBatchTaskSucceed = false;
	LOG.error("interrupted in executing batch-msg: " + _message.getId() + ", sub-msg: "
	+ subTask.getTaskId(), e);
	} catch (ExecutionException e) {
	isBatchTaskSucceed = false;
	LOG.error(
	"fail to execute batch-msg: " + _message.getId() + ", sub-msg: "
	+ subTask.getTaskId(), e);
	}
	}
	result = new HelixTaskResult();
	result.setSuccess(isBatchTaskSucceed);
	}

	result.setCompleteTime(System.currentTimeMillis());
	// pass task-result to post-handle-msg
	_notificationContext.add(MapKey.HELIX_TASK_RESULT.toString(), result);
	postHandleMessage();

	return result;
	}
	}

	@Override
	public void onError(Exception e, ErrorCode code, ErrorType type) {
	// if one sub-message execution fails, call onError on all sub-message handlers
	for (MessageHandler handler : _subMessageHandlers) {
	handler.onError(e, code, type);
	}
	}

	// TODO: optimize this based on the fact that each cs update is for a
	// distinct partition
	private Map<PropertyKey, CurrentState> mergeCurStateUpdate(
	ConcurrentHashMap<String, CurrentStateUpdate> csUpdateMap) {
	Map<String, CurrentStateUpdate> curStateUpdateMap = new HashMap<String, CurrentStateUpdate>();
	for (CurrentStateUpdate update : csUpdateMap.values()) {
	String path = update._key.getPath(); // TODO: this is time
	// consuming, optimize it
	if (!curStateUpdateMap.containsKey(path)) {
	curStateUpdateMap.put(path, update);
	} else {
	// long start = System.currentTimeMillis();
	curStateUpdateMap.get(path).merge(update._delta);
	// long end = System.currentTimeMillis();
	// LOG.info("each merge took: " + (end - start));
	}
	}

	Map<PropertyKey, CurrentState> ret = new HashMap<PropertyKey, CurrentState>();
	for (CurrentStateUpdate update : curStateUpdateMap.values()) {
	ret.put(update._key, update._delta);
	}

	return ret;
	}

	}