package org.apache.helix.manager.zk;
/*
* 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.io.StringReader;
import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.helix.Criteria;
import org.apache.helix.HelixDataAccessor;
import org.apache.helix.HelixException;
import org.apache.helix.HelixManager;
import org.apache.helix.NotificationContext;
import org.apache.helix.ZNRecord;
import org.apache.helix.PropertyKey.Builder;
import org.apache.helix.messaging.AsyncCallback;
import org.apache.helix.messaging.handling.HelixTaskResult;
import org.apache.helix.messaging.handling.MessageHandler;
import org.apache.helix.messaging.handling.MessageHandlerFactory;
import org.apache.helix.model.Message;
import org.apache.helix.model.StatusUpdate;
import org.apache.helix.model.Message.MessageType;
import org.apache.helix.util.StatusUpdateUtil;
import org.apache.log4j.Logger;
import org.codehaus.jackson.map.ObjectMapper;
/*
* TODO: The current implementation is temporary for backup handler testing only and it does not
* do any throttling.
*
*/
public class DefaultSchedulerMessageHandlerFactory implements
MessageHandlerFactory
{
public static final String WAIT_ALL = "WAIT_ALL";
public static final String SCHEDULER_MSG_ID = "SchedulerMessageId";
public static class SchedulerAsyncCallback extends AsyncCallback
{
StatusUpdateUtil _statusUpdateUtil = new StatusUpdateUtil();
Message _originalMessage;
HelixManager _manager;
final Map<String, Map<String, String>> _resultSummaryMap = new ConcurrentHashMap<String, Map<String, String>>();
public SchedulerAsyncCallback(Message originalMessage, HelixManager manager)
{
_originalMessage = originalMessage;
_manager = manager;
}
@Override
public void onTimeOut()
{
_logger.info("Scheduler msg timeout " + _originalMessage.getMsgId()
+ " timout with " + _timeout + " Ms");
_statusUpdateUtil.logError(_originalMessage,
SchedulerAsyncCallback.class, "Task timeout",
_manager.getHelixDataAccessor());
addSummary(_resultSummaryMap, _originalMessage, _manager, true);
}
@Override
public void onReplyMessage(Message message)
{
_logger.info("Update for scheduler msg " + _originalMessage.getMsgId()
+ " Message " + message.getMsgSrc() + " id "
+ message.getCorrelationId() + " completed");
String key = "MessageResult " + message.getMsgSrc() + " "
+ UUID.randomUUID();
_resultSummaryMap.put(key, message.getResultMap());
if (this.isDone())
{
_logger.info("Scheduler msg " + _originalMessage.getMsgId()
+ " completed");
_statusUpdateUtil.logInfo(_originalMessage,
SchedulerAsyncCallback.class, "Scheduler task completed",
_manager.getHelixDataAccessor());
addSummary(_resultSummaryMap, _originalMessage, _manager, false);
}
}
private void addSummary(Map<String, Map<String, String>> _resultSummaryMap,
Message originalMessage, HelixManager manager, boolean timeOut)
{
Map<String, String> summary = new TreeMap<String, String>();
summary.put("TotalMessages:", "" + _resultSummaryMap.size());
summary.put("Timeout", "" + timeOut);
_resultSummaryMap.put("Summary", summary);
HelixDataAccessor accessor = manager.getHelixDataAccessor();
Builder keyBuilder = accessor.keyBuilder();
ZNRecord statusUpdate = accessor.getProperty(
keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.toString(),
originalMessage.getMsgId())).getRecord();
statusUpdate.getMapFields().putAll(_resultSummaryMap);
accessor.setProperty(keyBuilder.controllerTaskStatus(
MessageType.SCHEDULER_MSG.toString(), originalMessage.getMsgId()),
new StatusUpdate(statusUpdate));
}
}
private static Logger _logger = Logger
.getLogger(DefaultSchedulerMessageHandlerFactory.class);
HelixManager _manager;
public DefaultSchedulerMessageHandlerFactory(HelixManager manager)
{
_manager = manager;
}
@Override
public MessageHandler createHandler(Message message,
NotificationContext context)
{
String type = message.getMsgType();
if (!type.equals(getMessageType()))
{
throw new HelixException("Unexpected msg type for message "
+ message.getMsgId() + " type:" + message.getMsgType());
}
return new DefaultSchedulerMessageHandler(message, context, _manager);
}
@Override
public String getMessageType()
{
return MessageType.SCHEDULER_MSG.toString();
}
@Override
public void reset()
{
}
public static class DefaultSchedulerMessageHandler extends MessageHandler
{
HelixManager _manager;
public DefaultSchedulerMessageHandler(Message message,
NotificationContext context, HelixManager manager)
{
super(message, context);
_manager = manager;
}
@Override
public HelixTaskResult handleMessage() throws InterruptedException
{
String type = _message.getMsgType();
HelixTaskResult result = new HelixTaskResult();
if (!type.equals(MessageType.SCHEDULER_MSG.toString()))
{
throw new HelixException("Unexpected msg type for message "
+ _message.getMsgId() + " type:" + _message.getMsgType());
}
// Parse timeout value
int timeOut = -1;
if (_message.getRecord().getSimpleFields().containsKey("TIMEOUT"))
{
try
{
timeOut = Integer.parseInt(_message.getRecord().getSimpleFields()
.get("TIMEOUT"));
} catch (Exception e)
{
}
}
// Parse the message template
ZNRecord record = new ZNRecord("templateMessage");
record.getSimpleFields().putAll(
_message.getRecord().getMapField("MessageTemplate"));
Message messageTemplate = new Message(record);
// Parse the criteria
StringReader sr = new StringReader(_message.getRecord().getSimpleField(
"Criteria"));
ObjectMapper mapper = new ObjectMapper();
Criteria recipientCriteria;
try
{
recipientCriteria = mapper.readValue(sr, Criteria.class);
} catch (Exception e)
{
_logger.error("", e);
result.setException(e);
result.setSuccess(false);
return result;
}
_logger.info("Scheduler sending message, criteria:" + recipientCriteria);
boolean waitAll = false;
if(_message.getRecord().getSimpleField(DefaultSchedulerMessageHandlerFactory.WAIT_ALL) !=null)
{
try
{
waitAll = Boolean.parseBoolean(_message.getRecord().getSimpleField(DefaultSchedulerMessageHandlerFactory.WAIT_ALL));
}
catch(Exception e)
{
_logger.warn("",e);
}
}
// Send all messages.
int nMsgsSent = 0;
SchedulerAsyncCallback callback = new SchedulerAsyncCallback(_message, _manager);
if(waitAll)
{
nMsgsSent = _manager.getMessagingService().sendAndWait(recipientCriteria,
messageTemplate,
callback,
timeOut);
}
else
{
nMsgsSent = _manager.getMessagingService().send(recipientCriteria,
messageTemplate,
callback,
timeOut);
}
HelixDataAccessor accessor = _manager.getHelixDataAccessor();
Builder keyBuilder = accessor.keyBuilder();
// Record the number of messages sent into status updates
Map<String, String> sendSummary = new HashMap<String, String>();
sendSummary.put("MessageCount", "" + nMsgsSent);
ZNRecord statusUpdate = accessor.getProperty(
keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.toString(),
_message.getMsgId())).getRecord();
statusUpdate.getMapFields().put("SentMessageCount", sendSummary);
accessor.setProperty(keyBuilder.controllerTaskStatus(
MessageType.SCHEDULER_MSG.toString(), _message.getMsgId()),
new StatusUpdate(statusUpdate));
result.getTaskResultMap().put(
"ControllerResult",
"msg " + _message.getMsgId() + " from " + _message.getMsgSrc()
+ " processed");
result.getTaskResultMap().put(SCHEDULER_MSG_ID, _message.getMsgId());
result.setSuccess(true);
return result;
}
@Override
public void onError(Exception e, ErrorCode code, ErrorType type)
{
_logger.error("Message handling pipeline get an exception. MsgId:"
+ _message.getMsgId(), e);
}
}
}