content: /**
   * 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 com.alibaba.jstorm.batch.meta;
  
  import java.util.ArrayList;
  import java.util.Date;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  
  import org.apache.log4j.Logger;
  
  import backtype.storm.topology.FailedException;
  
  import com.alibaba.jstorm.batch.BatchId;
  import com.alibaba.jstorm.batch.ICommitter;
  import com.alibaba.jstorm.batch.util.BatchDef;
  import com.alibaba.jstorm.cluster.ClusterState;
  import com.alibaba.jstorm.utils.JStormUtils;
  import com.alibaba.rocketmq.client.consumer.PullResult;
  import com.alibaba.rocketmq.client.consumer.PullStatus;
  import com.alibaba.rocketmq.client.exception.MQClientException;
  import com.alibaba.rocketmq.common.message.MessageExt;
  import com.alibaba.rocketmq.common.message.MessageQueue;
  import com.taobao.metaq.client.MetaPullConsumer;
  
  public class MetaSimpleClient implements ICommitter{
  	private static final Logger LOG = Logger.getLogger(MetaSimpleClient.class);
  	
  	private final MetaSpoutConfig metaSpoutConfig;
  	
  	private final int taskIndex;
  	private final int taskParallel;
  	private int       oneQueueFetchSize;
  	
  	private Map<MessageQueue, Long> currentOffsets;
  	private Map<MessageQueue, Long> frontOffsets;
  	private Map<MessageQueue, Long> backendOffset;
  	
  	private MetaPullConsumer        consumer;
  	
  	private static String           nameServer;
  	
  	private final ClusterState      zkClient;
  	
  	public MetaSimpleClient(MetaSpoutConfig config, 
  			ClusterState zkClient,
  			int taskIndex, 
  			int taskParall) {
  		this.metaSpoutConfig = config;
  		this.zkClient = zkClient;
  		this.taskIndex = taskIndex;
  		this.taskParallel = taskParall;
  	}
  	
  	public void initMetaConsumer() throws MQClientException {
  		LOG.info("MetaSpoutConfig:" + metaSpoutConfig);
  		
  		consumer = new MetaPullConsumer(metaSpoutConfig.getConsumerGroup());
  		consumer.setInstanceName(taskIndex + "." + JStormUtils.process_pid());
  		
  		if (metaSpoutConfig.getNameServer() != null) {
  			// this is for alipay
  			if (nameServer == null) {
  				nameServer = metaSpoutConfig.getNameServer();
  				
  				System.setProperty("rocketmq.namesrv.domain",
  						metaSpoutConfig.getNameServer());
  			}else if (metaSpoutConfig.getNameServer().equals(nameServer) == false) {
  				throw new RuntimeException("Different nameserver address in the same worker "
  						+ nameServer + ":" + metaSpoutConfig.getNameServer());
  				
  			}
  		}
  
          consumer.start();
          
          LOG.info("Successfully start meta consumer");
  	}
  	
  	public int getOneQueueFetchSize() {
  		int batchSize = metaSpoutConfig.getBatchMsgNum();
  		
  		int oneFetchSize = batchSize/(taskParallel * currentOffsets.size());
  		if (oneFetchSize <= 0) {
  			oneFetchSize = 1;
  		}
  		
  		LOG.info("One queue fetch size:" + oneFetchSize);
  		return oneFetchSize;
  	}
  	
  	public void init() throws Exception {
  		initMetaConsumer();
  		
  		frontOffsets = initOffset();
  		currentOffsets = frontOffsets;
  		backendOffset = new HashMap<MessageQueue, Long>();
  		backendOffset.putAll(frontOffsets);
  		
  		oneQueueFetchSize = getOneQueueFetchSize();
  	}
  	
  	
  	protected Set<MessageQueue> getMQ() throws MQClientException {
  		Set<MessageQueue> mqs = consumer.
  				fetchSubscribeMessageQueues(metaSpoutConfig.getTopic());
  		if (taskParallel > mqs.size()) {
  			throw new RuntimeException("Two much task to fetch " + metaSpoutConfig);
  		}
  		
  		List<MessageQueue> mqList = JStormUtils.mk_list(mqs);
  		Set<MessageQueue> ret = new HashSet<MessageQueue>();
  		for (int i = taskIndex; i < mqList.size(); i += taskParallel) {
  			ret.add(mqList.get(i));
  		}
  		
  		if (ret.size() == 0) {
  			throw new RuntimeException("No meta queue need to be consume");
  		}
  		return ret;
  	}
  	
  	protected Map<MessageQueue, Long> initOffset() throws MQClientException {
  		Set<MessageQueue> queues = getMQ();
  		
  		Map<MessageQueue, Long> ret = new HashMap<MessageQueue, Long>();
  		Set<MessageQueue>       noOffsetQueues = new HashSet<MessageQueue>();
  		
  		
  		if (metaSpoutConfig.getStartTimeStamp() != null) {
  			Long timeStamp = metaSpoutConfig.getStartTimeStamp();
  			
  			
  			for (MessageQueue mq : queues) {
  				long offset = consumer.searchOffset(mq, timeStamp);
                  if (offset >= 0) {
                  	LOG.info("Successfully get " + mq + " offset of timestamp " + new Date(timeStamp));
                  	ret.put(mq, offset);
                  }else {
                  	LOG.info("Failed to get " + mq + " offset of timestamp " + new Date(timeStamp));
                  	noOffsetQueues.add(mq);
                  }
  			}
  		}else  {
  			noOffsetQueues.addAll(queues);
  		}
  		
  		if (noOffsetQueues.size() == 0) {
  			return ret;
  		}
  		
  		for (MessageQueue mq : noOffsetQueues) {
  			long offset = getOffsetFromZk(mq);
  			
  			ret.put(mq, offset);
  		}
  		
  		return ret;
  	}
  	
  	
  	protected String getZkPath(MessageQueue mq) {
  		StringBuffer sb = new StringBuffer();
  		sb.append(BatchDef.ZK_SEPERATOR);
  		sb.append(metaSpoutConfig.getConsumerGroup());
  		sb.append(BatchDef.ZK_SEPERATOR);
  		sb.append(mq.getBrokerName());
  		sb.append(BatchDef.ZK_SEPERATOR);
  		sb.append(mq.getQueueId());
  		
  		return sb.toString();
  	}
  	
  	protected long getOffsetFromZk(MessageQueue mq) {
  		String path = getZkPath(mq);
  		
  		try {
  			if (zkClient.node_existed(path, false) == false) {
  				LOG.info("No zk node of " + path);
  				return 0;
  			}
  			
  			byte[] data = zkClient.get_data(path, false);
  			String value = new String(data);
  			
  			long ret = Long.valueOf(value);
  			return ret;
  		}catch (Exception e) {
  			LOG.warn("Failed to get offset,", e);
  			return 0;
  		}
  	}
  	
  	protected void updateOffsetToZk(MessageQueue mq, Long offset) throws Exception {
  		String path = getZkPath(mq);
  		byte[] data = String.valueOf(offset).getBytes();
  		zkClient.set_data(path, data);
  
  	}
  	
  	protected void  updateOffsetToZk(Map<MessageQueue, Long> mqs) throws Exception{
  		for (Entry<MessageQueue, Long> entry : mqs.entrySet()) {
  			MessageQueue mq = entry.getKey();
  			Long offset = entry.getValue();
  			
  			updateOffsetToZk(mq, offset);
  		}
  		
  		LOG.info("Update zk offset," + mqs);
  	}
  	
  	protected void switchOffsetMap() {
  		Map<MessageQueue, Long> tmp = frontOffsets;
  		
  		frontOffsets = backendOffset;
  		backendOffset = tmp;
  		
  		currentOffsets = frontOffsets;
  	}
  
  	@Override
  	public byte[] commit(BatchId id) throws FailedException {
  		try {
  			updateOffsetToZk(currentOffsets);
  			switchOffsetMap();
  		}catch(Exception e) {
  			LOG.warn("Failed to update offset to ZK", e);
  			throw new FailedException(e);
  		}
  		
  		return null;
  	}
  
  	@Override
  	public void revert(BatchId id, byte[] commitResult) {
  		try {
  			switchOffsetMap();
  			updateOffsetToZk(currentOffsets);
  			
  		}catch(Exception e) {
  			LOG.warn("Failed to update offset to ZK", e);
  			throw new FailedException(e);
  		}
  	}
  	
  	/**
  	 * rebalanceMqList must run after commit
  	 * 
  	 * @throws MQClientException
  	 */
  	public void rebalanceMqList() throws Exception {
  		LOG.info("Begin to do rebalance operation");
  		Set<MessageQueue> newMqs = getMQ();
  		
  		Set<MessageQueue> oldMqs = currentOffsets.keySet();
  		
  		if (oldMqs.equals(newMqs) == true) {
  			LOG.info("No change of meta queues " + newMqs);
  			return ;
  		}
  		
  		Set<MessageQueue> removeMqs = new HashSet<MessageQueue>();
  		removeMqs.addAll(oldMqs);
  		removeMqs.removeAll(newMqs);
  		
  		Set<MessageQueue> addMqs = new HashSet<MessageQueue>();
  		addMqs.addAll(newMqs);
  		addMqs.removeAll(oldMqs);
  		
  		LOG.info("Remove " + removeMqs);
  		for (MessageQueue mq : removeMqs) {
  			Long offset = frontOffsets.remove(mq);
  			updateOffsetToZk(mq, offset);
  			
  			backendOffset.remove(mq);
  			
  		}
  		
  		LOG.info("Add " + addMqs);
  		for (MessageQueue mq : addMqs) {
  			long offset = getOffsetFromZk(mq);
  			frontOffsets.put(mq, offset);
  			backendOffset.put(mq, offset);
  		}
  	}
  	
  	public List<MessageExt> fetchOneBatch() {
  		List<MessageExt> ret = new ArrayList<MessageExt>();
  		
  		
  		String subexpress = metaSpoutConfig.getSubExpress();
  		for(Entry<MessageQueue, Long>entry : currentOffsets.entrySet()) {
  			MessageQueue mq = entry.getKey();
  			Long offset = entry.getValue();
  			
  			
  			int fetchSize = 0;
  			int oneFetchSize = Math.min(oneQueueFetchSize, 32);
  			
  			while(fetchSize < oneQueueFetchSize) {
  				
  				PullResult pullResult = null;
  	            try {
  	                pullResult = consumer.pullBlockIfNotFound(mq, subexpress, offset, oneFetchSize);
  	                offset = pullResult.getNextBeginOffset();
  	                PullStatus status = pullResult.getPullStatus();
  	                if  (status == PullStatus.FOUND) {
  	                	List<MessageExt> msgList = pullResult.getMsgFoundList();
  	                    ret.addAll(msgList);
  	                    fetchSize += msgList.size();
  	                    continue;
  	                }else if (status ==  PullStatus.NO_MATCHED_MSG) {
  	                    continue;
  	                }else if (status == PullStatus.NO_NEW_MSG ) {
  	                    break;
  	                }else if (status == PullStatus.OFFSET_ILLEGAL) {
  	                    break;
  	                }else {
  	                
  	                    break;
  	                }
  	            }
  	            catch (Exception e) {
  	                LOG.warn("Failed to fetch messages of " + mq + ":" + offset, e);
  	                break;
  	            }
  			}
  			
  			backendOffset.put(mq, offset);
  		}
  		
  		return ret;
  	}
  
  	public void cleanup() {
  		consumer.shutdown();
  	}
  }
  