消息缓冲器

package yz.his.service;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import org.bson.Document;

import io.netty.util.internal.StringUtil;
import yb.common.util.AtomLock;
import yb.common.util.MD5;
import yb.common.util.ObjectTransfer;
import yb.common.util.ScheduledThreadPool;
import yb.common.util.Tester;
import yb.core.docdb.DocDB;
import yb.core.redis.RedisMap;
import yb.im.server.entity.UserMessageEntity;

public class MessageBuffer {
	private static Buffer buffer = new Buffer();
	private static MessageIO messageIO = new MessageIO();

	private MessageBuffer() {

	}

	/**
	 * 存入消息
	 * 
	 * @param message
	 */
	public static void put(UserMessageEntity message) {
		if (message == null)
			return;

		// 存入消息到redis
		messageIO.putMessage(message);

		// 存入未读数
		messageIO.unreadIncr(message);

		// 存入buffer
		buffer.put(message);
	}

	/**
	 * 通过chatId等参数获取一定数量的消息，优先从redis中获取
	 * 
	 * @param chatId
	 * @param currentUserId
	 * @param before
	 *            可以为null
	 * @param count
	 *            可以为null
	 * @return
	 * @throws Exception
	 */
	public static List<UserMessageEntity> get(String chatId, String currentUserId, Long before, Integer count)
			throws Exception {
		if (StringUtil.isNullOrEmpty(chatId) || StringUtil.isNullOrEmpty(currentUserId)
				|| (count != null && count <= 0))
			return null;
		return messageIO.getMessages(chatId, currentUserId, before, count);
	}

	/**
	 * 获取未读数
	 * 
	 * @param chatId
	 * @param currentUserId
	 * @return
	 */
	public static Long getUnread(String chatId, String currentUserId) {
		if (StringUtil.isNullOrEmpty(chatId) || StringUtil.isNullOrEmpty(currentUserId))
			return 0L;
		return messageIO.getUnread(chatId, currentUserId);
	}

	/**
	 * 设置已读
	 * 
	 * @param chatId
	 * @param currentUserId
	 * @return
	 */
	public static boolean setRead(String chatId, String currentUserId) {
		return messageIO.setRead(chatId, currentUserId);
	}

	/**
	 * 开启定时同步任务
	 */
	public static void startBufferSchedule() {
		BufferJob.start();
	}

	static class MessageIO {
		private final Integer Default_Count = 10;
		private final int Default_DataBase = 7;
		private final int Timeout = 24 * 60 * 60;

		private List<UserMessageEntity> chatMessages = null;
		private List<UserMessageEntity> chatMessagesInRedis = null;
		private List<UserMessageEntity> chatMessagesInMongo = null;

		private Map<String, UserMessageEntity> syncMessages = null;

		private RedisMap<String, UserMessageEntity> getChatRedisMap(String chatId) {
			if (chatId == null)
				return null;
			return new RedisMap<String, UserMessageEntity>(chatId, Timeout, null, Default_DataBase);
		}

		/**
		 * 存入消息
		 * 
		 * @param message
		 * @return
		 */
		public void putMessage(UserMessageEntity message) {
			if (message == null)
				return;
			String chatId = message.getChatId();
			String messageId = message.getMessageId();
			if (StringUtil.isNullOrEmpty(chatId) || StringUtil.isNullOrEmpty(messageId))
				return;
			RedisMap<String, UserMessageEntity> chatRedis = getChatRedisMap(chatId);
			if (chatRedis == null)
				return;
			chatRedis.put(message.getMessageId(), message);
		}

		/**
		 * 获取消息
		 * 
		 * @param chatId
		 * @param currentUserId
		 * @param before
		 *            可以为空
		 * @param count
		 *            可以为空，若为空则为默认条数
		 * @return
		 * @throws Exception
		 */
		public List<UserMessageEntity> getMessages(String chatId, String currentUserId, Long before, Integer count)
				throws Exception {
			if (count == null)
				count = Default_Count;
			if (chatMessages == null)
				chatMessages = new ArrayList<UserMessageEntity>();
			chatMessages.clear();

			getMessageFromRedis(chatId, currentUserId, before, count);

			int size = chatMessages.size();
			if (size < count) {
				if (size > 0)
					before = chatMessages.get(size - 1).getChatTime();
				getMessageFromMongoDB(chatId, currentUserId, before, count - size);
			}

			// 排序处理

			return chatMessages;
		}

		/**
		 * 从Redis中获取消息
		 * 
		 * @param chatId
		 * @param currentUserId
		 * @param before
		 * @param count
		 */
		private void getMessageFromRedis(String chatId, String currentUserId, Long before, Integer count) {
			RedisMap<String, UserMessageEntity> chatRM = getChatRedisMap(chatId);
			if (chatRM == null)
				return;
			if (chatMessagesInRedis == null)
				chatMessagesInRedis = new ArrayList<>();
			chatMessagesInRedis.clear();
			Collection<UserMessageEntity> values = chatRM.values();
			if (values == null || values.isEmpty())
				return;

			chatMessagesInRedis.addAll(values);
			java.util.Collections.sort(chatMessagesInRedis, new Comparator<UserMessageEntity>() {
				@Override
				public int compare(UserMessageEntity o1, UserMessageEntity o2) {
					return ObjectTransfer.intValue(o2.getChatTime() - o1.getChatTime());
				}
			});

			int rCount = 0;
			if (before == null || before.longValue() <= 0L) {
				// 直接取前count条符合条件的数据
				for (UserMessageEntity entity : chatMessagesInRedis) {
					if (currentUserId.equals(entity.getFromUserId()) || currentUserId.equals(entity.getToUserId())) {
						chatMessages.add(entity);
						rCount++;
					} else {
						continue;
					}
					if (rCount == count)
						break;
				}
			} else {
				// 若没有符合时间的要求
				if (chatMessagesInRedis.get(chatMessagesInRedis.size() - 1).getChatTime() > before)
					return;
				for (UserMessageEntity entity : chatMessagesInRedis) {
					if (entity.getChatTime() <= before && (currentUserId.equals(entity.getFromUserId())
							|| currentUserId.equals(entity.getToUserId()))) {
						chatMessages.add(entity);
						rCount++;
					} else {
						continue;
					}
					if (rCount == count)
						break;
				}
			}

		}

		/**
		 * 从mongDB中获取消息
		 * 
		 * @param chatId
		 * @param currentUserId
		 * @param before
		 * @param limit
		 * @throws Exception
		 */
		private void getMessageFromMongoDB(String chatId, String currentUserId, Long before, Integer limit)
				throws Exception {
			DocDB mongo = DocDB.mongo();
			Document query = new Document();
			query.put("ChatId", chatId);
			List<Document> or = new ArrayList<>();
			or.add(new Document("FromUserId", currentUserId));
			or.add(new Document("ToUserId", currentUserId));
			query.append("$or", or);
			if (before != null && before > 0)
				query.append("ChatTime", new Document("$lt", before));
			chatMessagesInMongo = mongo.mongoQuery(UserMessageEntity.class, query, new Document("ChatTime", -1), null,
					limit);

			if (chatMessagesInMongo != null && !chatMessagesInMongo.isEmpty())
				chatMessages.addAll(chatMessagesInMongo);
		}

		/**
		 * 将Redis中的消息同步到MongoDb中
		 * 
		 * @param syncMap
		 * @throws Exception
		 */
		private void redis2MongoDB(Map<String, List<String>> syncMap) throws Exception {
			if (syncMap == null || syncMap.isEmpty())
				return;
			if (syncMessages == null)
				syncMessages = new HashMap<String, UserMessageEntity>();
			syncMessages.clear();

			for (Entry<String, List<String>> entry : syncMap.entrySet()) {
				String chatId = entry.getKey();
				if (StringUtil.isNullOrEmpty(chatId))
					continue;
				List<String> value = entry.getValue();
				if (value != null && !value.isEmpty()) {
					RedisMap<String, UserMessageEntity> chatRedisMap = getChatRedisMap(chatId);
					if (chatRedisMap == null || chatRedisMap.isEmpty()) {
						if (chatRedisMap != null)
							chatRedisMap.clear();
						continue;
					}
					for (String msgId : value) {
						UserMessageEntity entity = chatRedisMap.get(msgId);
						if (entity == null)
							continue;
						syncMessages.put(msgId, entity);
					}
				}

			}

			if (syncMessages.isEmpty())
				return;

			Set<String> keySet = syncMessages.keySet();
			Document query = new Document();
			query.append("MessageId", DocDB.in(keySet));
			DocDB mongo = DocDB.mongo();
			List<String> msgIdsInMongo = mongo.mongoQueryOneCol(UserMessageEntity.class, query, String.class,
					"MessageId", null, null, null);
			if (!msgIdsInMongo.isEmpty()) {
				for (String key : msgIdsInMongo) {
					syncMessages.remove(key);
				}
			}

			// insert into mongdb
			if (!syncMessages.isEmpty())
				mongo.mongoInsertList(new ArrayList<UserMessageEntity>(syncMessages.values()));

			// 将redis中的messages移除
			for (Entry<String, List<String>> entry : syncMap.entrySet()) {
				String chatId = entry.getKey();
				List<String> list = entry.getValue();
				if (StringUtil.isNullOrEmpty(chatId) || list == null || list.isEmpty())
					continue;
				RedisMap<String, UserMessageEntity> chatRedisMap = getChatRedisMap(chatId);
				if (chatRedisMap.isEmpty()) {
					chatRedisMap.clear();
					continue;
				} else {
					for (String key : list) {
						chatRedisMap.remove(key);
					}
				}
			}

		}

		// 已读/未读
		/**
		 * 存入消息未读自增1，key为chatId和toUserId的联合key
		 * 
		 * @param message
		 * @return 当前未读数
		 */
		public Long unreadIncr(UserMessageEntity message) {
			if (message == null)
				return 0L;
			String chatId = message.getChatId();
			String toUserId = message.getToUserId();
			if (StringUtil.isNullOrEmpty(chatId) || StringUtil.isNullOrEmpty(toUserId))
				return 0L;
			return AtomLock.incr(chatId + toUserId, 1, 60 * 60 * 24);// 暂时定为1天
		}

		/**
		 * 设置为已读
		 * 
		 * @param chatId
		 * @param currentUserId
		 * @return
		 */
		public boolean setRead(String chatId, String currentUserId) {
			if (StringUtil.isNullOrEmpty(chatId) || StringUtil.isNullOrEmpty(currentUserId)) {
				return false;
			}
			AtomLock.del(chatId + currentUserId);
			return ObjectTransfer.longValue(AtomLock.val(chatId + currentUserId)) == 0L;
		}

		/**
		 * 获取未读数
		 * 
		 * @param chatId
		 * @param currentUserId
		 * @return
		 */
		public Long getUnread(String chatId, String currentUserId) {
			if (StringUtil.isNullOrEmpty(chatId) || StringUtil.isNullOrEmpty(currentUserId)) {
				return 0L;
			}
			return AtomLock.val(chatId + currentUserId);
		}
	}

	static class Buffer {
		private final String BufferName_Prefix = MessageBuffer.class.getName() + "_buffer_";
		private final String Sync_AtomicKey = MessageBuffer.class.getName() + "_sync";
		private Integer Sync_Interval = 10; // seconds
		private final Integer Default_DataBase = 7;

		private final String lockKey = BufferName_Prefix + "lock";

		private RedisMap<String, String> currentBuffer = null;
		private RedisMap<String, String> syncBuffer = null;
		Map<String, List<String>> syncMap = null;

		// 基础方法
		private boolean getLock() {
			return AtomLock.incr(lockKey, 1, Sync_Interval) == 1L;
		}

		private void unLock(String lockKey) {
			AtomLock.del(lockKey);
		}

		private String createIncreName(String namePrefix) {
			return namePrefix + ObjectTransfer.stringValue(AtomLock.incr(namePrefix, 1, 60 * 60));// 1h
		}

		// buffer相关
		/**
		 * 创建一个新的buffer
		 * 
		 * @return
		 */
		private void create() {
			String createIncreName = createIncreName(BufferName_Prefix);
			currentBuffer = new RedisMap<String, String>(createIncreName, 1, TimeUnit.HOURS, Default_DataBase);
		}

		private RedisMap<String, String> get(Long id) {
			return new RedisMap<String, String>(BufferName_Prefix + id, 1, TimeUnit.HOURS, Default_DataBase);
		}

		/**
		 * 获取需要同步的buffer
		 * 
		 * @param id
		 *            用于同步统计的buffer 自增id
		 */
		private void getSyncBuffer(Long id) {
			syncBuffer = null;
			if (id != null && id == AtomLock.val(BufferName_Prefix))
				return;
			syncBuffer = get(id);

			if (syncBuffer != null && syncBuffer.isEmpty()) {
				syncBuffer.clear();
				getSyncBuffer(AtomLock.incr(Sync_AtomicKey, 1, 60 * 60));// 1h
			} else {
				System.out.println(syncBuffer.getId());
			}
		}

		/**
		 * 将message存入buffer
		 * 
		 * @param message
		 */
		public void put(UserMessageEntity message) {
			if (message == null)
				return;
			String messageId = message.getMessageId();
			if (StringUtil.isNullOrEmpty(messageId))
				return;
			if (currentBuffer == null)
				create();
			currentBuffer.put(messageId, message.getChatId());
		}

		/**
		 * 定时消息同步工作调用
		 * 
		 * @throws Exception
		 */
		public void sync(int interval) throws Exception {
			if (Sync_Interval != interval)
				Sync_Interval = interval;

			// 一同步就create a buffer
			create();
			// if(true)
			// return;
			if (currentBuffer == null)
				return;
			if (!DocDB.isSupported())
				return;

			while (AtomLock.val(Sync_AtomicKey) + 2 < AtomLock.val(BufferName_Prefix)) {
				getSyncBuffer(AtomLock.val(Sync_AtomicKey));
				if (syncBuffer == null)
					break;

				// 获取锁
				if (getLock()) {
					// 将buffer内容处理为syncMap
					if (syncMap == null)
						syncMap = new HashMap<>();
					syncMap.clear();
					Set<String> keySet = syncBuffer.keySet();
					for (String key : keySet) {
						String chatId = syncBuffer.get(key);
						if (StringUtil.isNullOrEmpty(chatId))
							continue;
						List<String> list = syncMap.get(chatId);
						if (list == null) {
							list = new ArrayList<>();
							list.add(key);
							syncMap.put(chatId, list);
						} else {
							list.add(key);
						}
					}

					// 同步消息
					if (!syncMap.isEmpty()) {
						messageIO.redis2MongoDB(syncMap);
					}

					// 清除同步buffer
					syncBuffer.clear();

					// 解锁
					unLock(lockKey);
				}
			}

		}

	}

	static class BufferJob implements Runnable {
		private static final Integer Delay = 2;

		@Override
		public void run() {
			try {
				buffer.sync(Delay);
			} catch (Exception e) {
				e.printStackTrace();
			} finally {
				ScheduledThreadPool.getInstance().schedule(new BufferJob(), Delay, TimeUnit.SECONDS);
			}

		}

		public static void start() {
			ScheduledThreadPool.getInstance().schedule(new BufferJob(), Delay, TimeUnit.SECONDS);
		}
	}

	public static void main(String[] args) {
		try {
			Tester.init();
			UserMessageEntity entity = new UserMessageEntity();
			MessageBuffer.startBufferSchedule();
			for (int i = 0; i < 10000; i++) {

				entity.setChatId(MD5.hash(MessageBuffer.class.getName()));
				entity.setChatTime(System.currentTimeMillis());
				entity.setMessageId("messageId_2_" + i);
				entity.setFromUserId("this");
				entity.setToUserId("you");
				entity.setMessage("test redis");
				MessageBuffer.put(entity);
			}
			System.out.println("已经存入消息");
			System.out.println("未读消息数："+MessageBuffer.getUnread(MD5.hash(MessageBuffer.class.getName()), "you"));
			MessageBuffer.setRead(MD5.hash(MessageBuffer.class.getName()), "you");
			System.out.println("未读消息数："+MessageBuffer.getUnread(MD5.hash(MessageBuffer.class.getName()), "you"));
			Thread.sleep(10000);

			List<UserMessageEntity> list = MessageBuffer.get(MD5.hash(MessageBuffer.class.getName()), "this", null,
					11002);
			System.out.println("获取到的消息列表：");
			System.out.println(ObjectTransfer.print(list.size()));

		} catch (Exception e) {
			Tester.rollback();
			e.printStackTrace();
		} finally {
			Tester.close();
		}
	}
}