All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.lumongo.server.index.LumongoIndexManager Maven / Gradle / Ivy

The newest version!
package org.lumongo.server.index;

import com.google.protobuf.util.JsonFormat;
import com.hazelcast.core.IExecutorService;
import com.hazelcast.core.Member;
import com.mongodb.MongoClient;
import com.mongodb.MongoException;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.MongoIterable;
import org.apache.log4j.Logger;
import org.apache.lucene.facet.FacetsConfig;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TermQuery;
import org.bson.Document;
import org.lumongo.cluster.message.Lumongo;
import org.lumongo.cluster.message.Lumongo.*;
import org.lumongo.server.config.ClusterConfig;
import org.lumongo.server.config.IndexConfig;
import org.lumongo.server.config.LocalNodeConfig;
import org.lumongo.server.config.MongoConfig;
import org.lumongo.server.config.Nodes;
import org.lumongo.server.connection.InternalClient;
import org.lumongo.server.connection.SocketRequestFederator;
import org.lumongo.server.exceptions.IndexDoesNotExist;
import org.lumongo.server.exceptions.InvalidIndexConfig;
import org.lumongo.server.hazelcast.HazelcastManager;
import org.lumongo.server.hazelcast.ReloadIndexSettingsTask;
import org.lumongo.server.hazelcast.UnloadIndexTask;
import org.lumongo.server.search.QueryCombiner;
import org.lumongo.server.search.QueryWithFilters;
import org.lumongo.util.ClusterHelper;
import org.lumongo.util.LumongoThreadFactory;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class LumongoIndexManager {
	private final static Logger log = Logger.getLogger(LumongoIndexManager.class);

	private final ReadWriteLock globalLock;

	private final ConcurrentHashMap indexMap;
	private final InternalClient internalClient;

	private final ExecutorService pool;
	private final ClusterHelper clusterHelper;

	private HazelcastManager hazelcastManager;

	private MongoConfig mongoConfig;
	private ClusterConfig clusterConfig;

	private MongoClient mongo;

	private AtomicLong queryNumber;

	public LumongoIndexManager(MongoClient mongo, MongoConfig mongoConfig, ClusterConfig clusterConfig) throws UnknownHostException {
		this.globalLock = new ReentrantReadWriteLock(true);

		this.mongoConfig = mongoConfig;
		this.clusterConfig = clusterConfig;

		this.indexMap = new ConcurrentHashMap<>();

		this.mongo = mongo;
		this.clusterHelper = new ClusterHelper(mongo, mongoConfig.getDatabaseName());
		this.internalClient = new InternalClient(clusterHelper, clusterConfig);

		this.pool = Executors.newCachedThreadPool(new LumongoThreadFactory("manager"));

		queryNumber = new AtomicLong();

	}

	public ClusterConfig getClusterConfig() {
		return clusterConfig;
	}

	public void init(HazelcastManager hazelcastManager) throws UnknownHostException, MongoException {
		globalLock.writeLock().lock();
		try {
			this.hazelcastManager = hazelcastManager;
		}
		finally {
			globalLock.writeLock().unlock();
		}

	}

	public void handleServerRemoved(Set currentMembers, Member memberRemoved, boolean master) {
		globalLock.writeLock().lock();
		try {
			if (master) {
				handleServerRemoved(currentMembers, memberRemoved);
			}

			internalClient.removeMember(memberRemoved);
		}
		finally {
			globalLock.writeLock().unlock();
		}

	}

	public void handleServerAdded(Set currentMembers, Member memberAdded, boolean master) throws Exception {
		globalLock.writeLock().lock();
		try {
			if (master) {
				// make sure we can resolve it before transferring segments
				Nodes nodes = clusterHelper.getNodes();
				@SuppressWarnings("unused") LocalNodeConfig localNodeConfig = nodes.find(memberAdded);

				handleServerAdded(currentMembers, memberAdded);
			}

			internalClient.addMember(memberAdded);
		}
		finally {
			globalLock.writeLock().unlock();
		}

	}

	public List getIndexNames() {
		globalLock.writeLock().lock();

		try {
			ArrayList indexNames = new ArrayList<>();
			log.info("Searching database <" + mongoConfig.getDatabaseName() + "> for indexes");
			MongoDatabase db = mongo.getDatabase(mongoConfig.getDatabaseName());
			MongoIterable allCollections = db.listCollectionNames();

			for (String collection : allCollections) {
				if (collection.endsWith(LumongoIndex.CONFIG_SUFFIX)) {
					String indexName = collection.substring(0, collection.length() - LumongoIndex.CONFIG_SUFFIX.length());
					indexNames.add(indexName);
				}

			}

			return indexNames;
		}
		finally {
			globalLock.writeLock().unlock();
		}
	}

	public void loadIndexes() {
		globalLock.writeLock().lock();

		try {
			log.info("Loading existing indexes");
			List indexNames = getIndexNames();
			for (String indexName : indexNames) {
				try {
					loadIndex(indexName, true);
				}
				catch (Exception e) {
					log.error("Failed to load index <" + indexName + ">: " + e.getClass().getSimpleName() + ": ", e);
				}
			}
			log.info("Finished loading existing indexes");
		}
		finally {
			globalLock.writeLock().unlock();
		}
	}

	public IndexCreateResponse createIndex(IndexCreateRequest request) throws Exception {
		globalLock.writeLock().lock();
		try {
			log.info("Creating index: <" + request.getIndexName() + ">:\n" + JsonFormat.printer().print(request));

			IndexConfig indexConfig = new IndexConfig(request);

			String indexName = indexConfig.getIndexName();
			if (indexMap.containsKey(indexName)) {
				throw new Exception("Index <" + indexName + "> already exist");
			}
			LumongoIndex i = LumongoIndex.createIndex(hazelcastManager, mongoConfig, clusterConfig, indexConfig);
			indexMap.put(indexConfig.getIndexName(), i);
			i.loadAllSegments();
			i.forceBalance(hazelcastManager.getMembers());

			log.info("Created index: <" + request.getIndexName() + ">");

			return IndexCreateResponse.newBuilder().build();
		}
		finally {
			globalLock.writeLock().unlock();
		}
	}

	public IndexConfig getIndexConfig(String indexName) throws InvalidIndexConfig {
		return LumongoIndex.loadIndexSettings(mongo, mongoConfig.getDatabaseName(), indexName);
	}

	public void loadIndex(String indexName, boolean loadAllSegments) throws Exception {
		globalLock.writeLock().lock();
		try {
			LumongoIndex i = LumongoIndex.loadIndex(hazelcastManager, mongoConfig, mongo, clusterConfig, indexName);
			if (loadAllSegments) {
				i.loadAllSegments();
			}
			indexMap.put(indexName, i);
		}
		finally {
			globalLock.writeLock().unlock();
		}
	}

	private void handleServerAdded(Set currentMembers, Member memberAdded) {
		globalLock.writeLock().lock();
		try {
			for (String key : indexMap.keySet()) {
				LumongoIndex i = indexMap.get(key);
				i.handleServerAdded(currentMembers, memberAdded);
			}
		}
		finally {
			globalLock.writeLock().unlock();
		}
	}

	private void handleServerRemoved(Set currentMembers, Member memberRemoved) {
		globalLock.writeLock().lock();
		try {
			for (String key : indexMap.keySet()) {
				LumongoIndex i = indexMap.get(key);
				i.handleServerRemoved(currentMembers, memberRemoved);
			}
		}
		finally {
			globalLock.writeLock().unlock();
		}
	}

	public void updateSegmentMap(String indexName, Map> newMemberToSegmentMap) throws Exception {
		globalLock.writeLock().lock();
		try {
			if (!indexMap.containsKey(indexName)) {
				loadIndex(indexName, false);
			}

			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}
			i.updateSegmentMap(newMemberToSegmentMap);
		}
		finally {
			globalLock.writeLock().unlock();
		}
	}

	public IndexDeleteResponse deleteIndex(IndexDeleteRequest request) throws Exception {
		globalLock.writeLock().lock();
		try {
			String indexName = request.getIndexName();

			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				if (getIndexNames().contains(indexName)) {
					//TODO delete index from database
					return IndexDeleteResponse.newBuilder().build();
				}
				else {
					throw new IndexDoesNotExist(indexName);
				}
			}

			Set currentMembers = hazelcastManager.getMembers();
			IExecutorService executorService = hazelcastManager.getExecutorService();

			Member self = hazelcastManager.getSelf();

			log.info("Unload index <" + indexName + "> for delete");
			for (Member m : currentMembers) {
				try {
					UnloadIndexTask uit = new UnloadIndexTask(m.getSocketAddress().getPort(), indexName, true);
					if (!self.equals(m)) {
						Future dt = executorService.submitToMember(uit, m);
						dt.get();
					}
					else {
						uit.call();
					}
				}
				catch (Exception e) {
					log.error(e.getClass().getSimpleName() + ": ", e);
				}

			}

			log.info("Deleting index <" + indexName + ">");
			i.deleteIndex();
			indexMap.remove(indexName);

			return IndexDeleteResponse.newBuilder().build();
		}
		finally {
			globalLock.writeLock().unlock();
		}
	}

	public void unloadIndex(String indexName, boolean terminate) throws IOException {
		globalLock.writeLock().lock();
		try {
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			i.unload(terminate);
			indexMap.remove(indexName);
		}
		finally {
			globalLock.writeLock().unlock();
		}
	}

	public void shutdown() {

		log.info("Starting index manager shutdown");

		//TODO configure or force a lock acquire
		int waitSeconds = 10;

		log.info("Waiting for lock");
		boolean locked = false;
		try {
			locked = globalLock.writeLock().tryLock(waitSeconds, TimeUnit.SECONDS);
		}
		catch (InterruptedException e1) {

		}

		if (!locked) {
			log.info("Failed to get manager lock within <" + waitSeconds + "> seconds");
		}

		try {
			log.info("Stopping manager pool");
			pool.shutdownNow();

			log.info("Shutting down indexes");
			for (String indexName : indexMap.keySet()) {
				LumongoIndex i = indexMap.get(indexName);
				try {
					log.info("Unloading <" + indexName + ">");
					i.unload(false);
				}
				catch (Exception e) {
					log.error(e.getClass().getSimpleName() + ": ", e);
				}
			}
		}
		finally {
			if (locked) {
				globalLock.writeLock().unlock();
			}
		}
	}

	public void openConnections(Set members) throws Exception {

		globalLock.writeLock().lock();
		try {
			Member self = hazelcastManager.getSelf();
			for (Member m : members) {
				if (!self.equals(m)) {
					internalClient.addMember(m);
				}
			}
		}
		finally {
			globalLock.writeLock().unlock();
		}

	}

	public IndexSettingsResponse updateIndex(String indexName, IndexSettings request) throws InvalidIndexConfig, MongoException, IOException {
		globalLock.readLock().lock();
		try {
			log.info("Updating index settings for <" + indexName + ">:\n" + JsonFormat.printer().print(request));
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			i.updateIndexSettings(request);

			Set currentMembers = hazelcastManager.getMembers();
			IExecutorService executorService = hazelcastManager.getExecutorService();

			Member self = hazelcastManager.getSelf();

			for (Member m : currentMembers) {
				try {
					ReloadIndexSettingsTask rist = new ReloadIndexSettingsTask(m.getSocketAddress().getPort(), indexName);
					if (!self.equals(m)) {
						Future dt = executorService.submitToMember(rist, m);
						dt.get();
					}
					else {
						rist.call();
					}
				}
				catch (Exception e) {
					log.error(e.getClass().getSimpleName() + ": ", e);
				}

			}

			return IndexSettingsResponse.newBuilder().build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public void reloadIndexSettings(String indexName) throws Exception {
		globalLock.readLock().lock();
		try {

			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			i.reloadIndexSettings();

		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public DeleteResponse internalDeleteDocument(DeleteRequest deleteRequest) throws Exception {
		globalLock.readLock().lock();
		try {

			LumongoIndex i = indexMap.get(deleteRequest.getIndexName());
			if (i == null) {
				throw new IndexDoesNotExist(deleteRequest.getIndexName());
			}
			i.deleteDocument(deleteRequest);

			return DeleteResponse.newBuilder().build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public FetchResponse internalFetch(FetchRequest fetchRequest) throws Exception {
		globalLock.readLock().lock();
		try {

			LumongoIndex i = indexMap.get(fetchRequest.getIndexName());
			if (i == null) {
				throw new IndexDoesNotExist(fetchRequest.getIndexName());
			}

			FetchResponse.Builder frBuilder = FetchResponse.newBuilder();

			String uniqueId = fetchRequest.getUniqueId();

			FetchType resultFetchType = fetchRequest.getResultFetchType();
			if (!FetchType.NONE.equals(resultFetchType)) {

				Long timestamp = null;
				if (fetchRequest.hasTimestamp()) {
					timestamp = fetchRequest.getTimestamp();
				}

				ResultDocument resultDoc = i.getSourceDocument(uniqueId, timestamp, resultFetchType, fetchRequest.getDocumentFieldsList(),
						fetchRequest.getDocumentMaskedFieldsList(), Collections.emptyList());
				if (null != resultDoc) {
					frBuilder.setResultDocument(resultDoc);
				}
			}

			FetchType associatedFetchType = fetchRequest.getAssociatedFetchType();
			if (!FetchType.NONE.equals(associatedFetchType)) {
				if (fetchRequest.hasFilename()) {
					AssociatedDocument ad = i.getAssociatedDocument(uniqueId, fetchRequest.getFilename(), associatedFetchType);
					if (ad != null) {
						frBuilder.addAssociatedDocument(ad);
					}
				}
				else {
					for (AssociatedDocument ad : i.getAssociatedDocuments(uniqueId, associatedFetchType)) {
						frBuilder.addAssociatedDocument(ad);
					}
				}
			}
			return frBuilder.build();

		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public DeleteResponse deleteDocument(DeleteRequest deleteRequest) throws Exception {
		globalLock.readLock().lock();
		try {

			String indexName = deleteRequest.getIndexName();
			String uniqueId = deleteRequest.getUniqueId();

			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			Member m = i.findMember(uniqueId);

			Member self = hazelcastManager.getSelf();

			if (!self.equals(m)) {
				return internalClient.executeDelete(m, deleteRequest);
			}
			else {
				return internalDeleteDocument(deleteRequest);
			}

		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public StoreResponse storeInternal(StoreRequest storeRequest) throws Exception {
		globalLock.readLock().lock();
		try {
			String indexName = storeRequest.getIndexName();
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			i.storeInternal(storeRequest);

			return StoreResponse.newBuilder().build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public StoreResponse storeDocument(StoreRequest storeRequest) throws Exception {
		globalLock.readLock().lock();
		try {

			String uniqueId = storeRequest.getUniqueId();
			String indexName = storeRequest.getIndexName();

			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			Member m = i.findMember(uniqueId);

			Member self = hazelcastManager.getSelf();

			if (!self.equals(m)) {
				return internalClient.executeStore(m, storeRequest);
			}
			else {
				return storeInternal(storeRequest);
			}

		}
		finally {
			globalLock.readLock().unlock();
		}

	}

	public FetchResponse fetch(FetchRequest request) throws Exception {
		globalLock.readLock().lock();
		try {

			String indexName = request.getIndexName();

			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			Member m = i.findMember(request.getUniqueId());

			Member self = hazelcastManager.getSelf();

			if (!self.equals(m)) {
				return internalClient.executeFetch(m, request);
			}
			else {
				return internalFetch(request);
			}

		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	private Map getQueryMap(QueryRequest queryRequest) throws Exception {
		globalLock.readLock().lock();
		try {

			List indexNames = queryRequest.getIndexList();

			HashMap queryMap = new HashMap<>();
			for (String indexName : indexNames) {
				LumongoIndex i = indexMap.get(indexName);
				if (i == null) {
					throw new IndexDoesNotExist(indexName);
				}

				Query query = i.getQuery(queryRequest.getQuery());

				QueryWithFilters queryWithFilters = new QueryWithFilters(query);

				for (FieldSimilarity fieldSimilarity : queryRequest.getFieldSimilarityList()) {
					queryWithFilters.addSimilarityOverride(fieldSimilarity);
				}



				if (queryRequest.hasFacetRequest()) {
					FacetRequest facetRequest = queryRequest.getFacetRequest();

					List drillDownList = facetRequest.getDrillDownList();
					if (!drillDownList.isEmpty()) {


						Map> dimToValues = new HashMap<>();
						for (LMFacet drillDown : drillDownList) {
							String key = drillDown.getLabel();
							String value = drillDown.getPath();
							if (!dimToValues.containsKey(key)) {
								dimToValues.put(key, new HashSet<>());
							}
							dimToValues.get(key).add(value);
						}

						for (Map.Entry> entry : dimToValues.entrySet()) {
							String indexFieldName = FacetsConfig.DEFAULT_INDEX_FIELD_NAME + "." + entry.getKey();

							BooleanQuery.Builder booleanQuery = new BooleanQuery.Builder();
							for (String value : entry.getValue()) {
								booleanQuery.add(new BooleanClause(new TermQuery(new org.apache.lucene.index.Term(indexFieldName, value)),
										BooleanClause.Occur.SHOULD));
							}

							queryWithFilters.addFilterQuery(booleanQuery.build());
						}

					}
				}

				for (Lumongo.Query filterQuery : queryRequest.getFilterQueryList()) {
					queryWithFilters.addFilterQuery(i.getQuery(filterQuery));
				}

				for (CosineSimRequest cosineSimRequest : queryRequest.getCosineSimRequestList()) {
					i.handleCosineSimQuery(queryWithFilters, cosineSimRequest);
				}

				queryMap.put(indexName, queryWithFilters);
			}

			return queryMap;
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public QueryResponse query(final QueryRequest request) throws Exception {
		globalLock.readLock().lock();
		long start = System.currentTimeMillis();
		long queryId = queryNumber.getAndIncrement();
		try {
			//log.info("Running query: <" + request.getQuery() + "> on indexes <" + request.getIndexList() + ">");

			String queryJson = JsonFormat.printer().print(request);
			log.info("Running id <" + queryId + "> query <" + queryJson + ">");

			final Map queryMap = getQueryMap(request);

			final Map indexSegmentMap = new HashMap<>();
			for (String indexName : request.getIndexList()) {
				LumongoIndex i = indexMap.get(indexName);
				if (i == null) {
					throw new IndexDoesNotExist(indexName);
				}
				indexSegmentMap.put(indexName, i);
			}

			SocketRequestFederator queryFederator = new SocketRequestFederator(
					hazelcastManager, pool) {

				@Override
				public InternalQueryResponse processExternal(Member m, QueryRequest request) throws Exception {
					return internalClient.executeQuery(m, request);
				}

				@Override
				public InternalQueryResponse processInternal(QueryRequest request) throws Exception {
					return internalQuery(queryMap, request);
				}
			};

			List results = queryFederator.send(request);

			QueryCombiner queryCombiner = new QueryCombiner(indexSegmentMap, request, results);

			queryCombiner.validate();

			QueryResponse qr = queryCombiner.getQueryResponse();

			if (!queryCombiner.isShort()) {
				return qr;
			}
			else {
				if (!request.getFetchFull()) {
					return query(request.toBuilder().setFetchFull(true).build());
				}

				throw new Exception("Full fetch request is short");
			}

		}
		finally {
			long end = System.currentTimeMillis();
			log.info("Finished query id <" + queryId + "> in " + (end - start) + "ms");

			globalLock.readLock().unlock();
		}
	}

	public InternalQueryResponse internalQuery(QueryRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			Map queryMap = getQueryMap(request);
			return internalQuery(queryMap, request);
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	private InternalQueryResponse internalQuery(Map queryMap, QueryRequest request) throws Exception {
		globalLock.readLock().lock();
		try {

			InternalQueryResponse.Builder internalQueryResponseBuilder = InternalQueryResponse.newBuilder();
			for (String indexName : queryMap.keySet()) {

				LumongoIndex i = indexMap.get(indexName);
				if (i == null) {
					throw new IndexDoesNotExist(indexName);
				}
				QueryWithFilters queryWithFilters = queryMap.get(indexName);

				IndexSegmentResponse isr = i.queryInternal(queryWithFilters, request);
				internalQueryResponseBuilder.addIndexSegmentResponse(isr);
			}

			return internalQueryResponseBuilder.build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public GetIndexesResponse getIndexes(GetIndexesRequest request) {
		globalLock.readLock().lock();
		try {
			GetIndexesResponse.Builder girB = GetIndexesResponse.newBuilder();
			girB.addAllIndexName(indexMap.keySet());
			return girB.build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public GetNumberOfDocsResponse getNumberOfDocsInternal(GetNumberOfDocsRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			String indexName = request.getIndexName();
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}
			return i.getNumberOfDocs();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public GetNumberOfDocsResponse getNumberOfDocs(GetNumberOfDocsRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			String indexName = request.getIndexName();
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}
			int numberOfSegments = i.getNumberOfSegments();

			SocketRequestFederator federator = new SocketRequestFederator(
					hazelcastManager, pool) {

				@Override
				public GetNumberOfDocsResponse processExternal(Member m, GetNumberOfDocsRequest request) throws Exception {
					return internalClient.getNumberOfDocs(m, request);
				}

				@Override
				public GetNumberOfDocsResponse processInternal(GetNumberOfDocsRequest request) throws Exception {
					return getNumberOfDocsInternal(request);
				}

			};

			GetNumberOfDocsResponse.Builder responseBuilder = GetNumberOfDocsResponse.newBuilder();
			responseBuilder.setNumberOfDocs(0);
			List responses = federator.send(request);

			List segmentCountResponses = new ArrayList<>();

			for (GetNumberOfDocsResponse r : responses) {
				responseBuilder.setNumberOfDocs(responseBuilder.getNumberOfDocs() + r.getNumberOfDocs());
				segmentCountResponses.addAll(r.getSegmentCountResponseList());
			}

			Collections.sort(segmentCountResponses, (o1, o2) -> Integer.compare(o1.getSegmentNumber(), o2.getSegmentNumber()));

			responseBuilder.addAllSegmentCountResponse(segmentCountResponses);

			GetNumberOfDocsResponse response = responseBuilder.build();

			HashSet segments = new HashSet<>();

			for (SegmentCountResponse r : response.getSegmentCountResponseList()) {
				segments.add(r.getSegmentNumber());
			}

			if (segments.size() != numberOfSegments) {
				throw new Exception("Expected <" + numberOfSegments + "> segments, found <" + segments.size() + "> segments");
			}

			for (int segmentNumber = 0; segmentNumber < numberOfSegments; segmentNumber++) {
				if (!segments.contains(segmentNumber)) {
					throw new Exception("Missing results for segment <" + segmentNumber + ">");
				}
			}

			return response;

		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public ClearResponse clearIndex(ClearRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			SocketRequestFederator federator = new SocketRequestFederator(hazelcastManager, pool) {

				@Override
				public ClearResponse processExternal(Member m, ClearRequest request) throws Exception {
					return internalClient.clear(m, request);
				}

				@Override
				public ClearResponse processInternal(ClearRequest request) throws Exception {
					return clearInternal(request);
				}

			};

			// nothing in responses currently
			@SuppressWarnings("unused") List responses = federator.send(request);

			return ClearResponse.newBuilder().build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public ClearResponse clearInternal(ClearRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			String indexName = request.getIndexName();
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}
			i.clear();
			return ClearResponse.newBuilder().build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public OptimizeResponse optimize(OptimizeRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			SocketRequestFederator federator = new SocketRequestFederator(
					hazelcastManager, pool) {

				@Override
				public OptimizeResponse processExternal(Member m, OptimizeRequest request) throws Exception {
					return internalClient.optimize(m, request);
				}

				@Override
				public OptimizeResponse processInternal(OptimizeRequest request) throws Exception {
					return optimizeInternal(request);
				}

			};

			// nothing in responses currently
			@SuppressWarnings("unused") List responses = federator.send(request);

			return OptimizeResponse.newBuilder().build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public OptimizeResponse optimizeInternal(OptimizeRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			String indexName = request.getIndexName();
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}
			i.optimize();
			return OptimizeResponse.newBuilder().build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public GetFieldNamesResponse getFieldNames(GetFieldNamesRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			SocketRequestFederator federator = new SocketRequestFederator(
					hazelcastManager, pool) {

				@Override
				public GetFieldNamesResponse processExternal(Member m, GetFieldNamesRequest request) throws Exception {
					return internalClient.getFieldNames(m, request);
				}

				@Override
				public GetFieldNamesResponse processInternal(GetFieldNamesRequest request) throws Exception {
					return getFieldNamesInternal(request);
				}

			};

			Set fieldNames = new HashSet<>();
			List responses = federator.send(request);
			for (GetFieldNamesResponse response : responses) {
				fieldNames.addAll(response.getFieldNameList());
			}

			GetFieldNamesResponse.Builder responseBuilder = GetFieldNamesResponse.newBuilder();
			responseBuilder.addAllFieldName(fieldNames);
			return responseBuilder.build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public GetFieldNamesResponse getFieldNamesInternal(GetFieldNamesRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			String indexName = request.getIndexName();
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}
			return i.getFieldNames();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public GetTermsResponse getTerms(GetTermsRequest request) throws Exception {

		globalLock.readLock().lock();
		try {

			SocketRequestFederator federator = new SocketRequestFederator(
					hazelcastManager, pool) {

				@Override
				public GetTermsResponseInternal processExternal(Member m, GetTermsRequest request) throws Exception {
					return internalClient.getTerms(m, request);
				}

				@Override
				public GetTermsResponseInternal processInternal(GetTermsRequest request) throws Exception {
					return getTermsInternal(request);
				}

			};

			List responses = federator.send(request);

			TreeMap terms = new TreeMap<>();
			for (GetTermsResponseInternal response : responses) {
				for (GetTermsResponse gtr : response.getGetTermsResponseList()) {
					for (Term term : gtr.getTermList()) {
						String key = term.getValue();
						if (!terms.containsKey(key)) {
							Term.Builder termBuilder = Term.newBuilder().setValue(key).setDocFreq(0).setTermFreq(0);
							if (term.hasScore()) {
								termBuilder.setScore(0);
							}
							terms.put(key, termBuilder);
						}
						Term.Builder builder = terms.get(key);
						builder.setDocFreq(builder.getDocFreq() + term.getDocFreq());
						builder.setTermFreq(builder.getTermFreq() + term.getTermFreq());
						if (term.hasScore()) {
							builder.setScore(builder.getScore() + term.getScore());
						}
					}
				}
			}

			GetTermsResponse.Builder responseBuilder = GetTermsResponse.newBuilder();

			Term.Builder value = null;

			int count = 0;

			int amount = request.getAmount();
			for (Term.Builder builder : terms.values()) {
				value = builder;
				if (builder.getDocFreq() >= request.getMinDocFreq() && builder.getTermFreq() >= request.getMinTermFreq()) {
					responseBuilder.addTerm(builder.build());
					count++;
				}

				if (amount != 0 && count >= amount) {
					break;
				}
			}

			if (value != null) {
				responseBuilder.setLastTerm(value.build());
			}

			return responseBuilder.build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public GetTermsResponseInternal getTermsInternal(GetTermsRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			String indexName = request.getIndexName();
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}
			return i.getTerms(request);
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	//rest
	public void storeAssociatedDocument(String indexName, String uniqueId, String fileName, InputStream is, boolean compress,
			HashMap metadataMap) throws Exception {
		globalLock.readLock().lock();
		try {
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			i.storeAssociatedDocument(uniqueId, fileName, is, compress, hazelcastManager.getClusterTime(), metadataMap);

		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public void getAssociatedDocuments(String indexName, OutputStream outputStream, Document filter) throws IOException {
		globalLock.readLock().lock();
		try {
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			i.getAssociatedDocuments(outputStream, filter);

		}
		finally {
			globalLock.readLock().unlock();
		}
	}


	//rest
	public InputStream getAssociatedDocumentStream(String indexName, String uniqueId, String fileName) throws IOException {
		globalLock.readLock().lock();
		try {
			LumongoIndex i = indexMap.get(indexName);
			if (i == null) {
				throw new IndexDoesNotExist(indexName);
			}

			return i.getAssociatedDocumentStream(uniqueId, fileName);

		}
		finally {
			globalLock.readLock().unlock();
		}
	}

	public GetMembersResponse getMembers(GetMembersRequest request) throws Exception {
		globalLock.readLock().lock();
		try {
			Set members = hazelcastManager.getMembers();
			GetMembersResponse.Builder responseBuilder = GetMembersResponse.newBuilder();

			Nodes nodes = clusterHelper.getNodes();

			HashMap memberMap = new HashMap<>();

			for (Member m : members) {
				LocalNodeConfig localNodeConfig = nodes.find(m);

				InetAddress inetAddress = m.getSocketAddress().getAddress();

				String fullHostName = inetAddress.getCanonicalHostName();

				LMMember.Builder lmMemberBuilder = LMMember.newBuilder();
				lmMemberBuilder.setServerAddress(fullHostName);
				lmMemberBuilder.setExternalPort(localNodeConfig.getExternalServicePort());
				lmMemberBuilder.setInternalPort(localNodeConfig.getInternalServicePort());
				lmMemberBuilder.setHazelcastPort(localNodeConfig.getHazelcastPort());
				lmMemberBuilder.setRestPort(localNodeConfig.getRestPort());
				LMMember lmMember = lmMemberBuilder.build();
				responseBuilder.addMember(lmMember);
				memberMap.put(m, lmMember);
			}

			for (String indexName : indexMap.keySet()) {
				LumongoIndex i = indexMap.get(indexName);

				IndexMapping.Builder indexMappingBuilder = IndexMapping.newBuilder();
				indexMappingBuilder.setIndexName(indexName);
				indexMappingBuilder.setNumberOfSegments(i.getNumberOfSegments());

				Map segmentToMemberMap = i.getSegmentToMemberMap();
				for (Integer segmentNumber : segmentToMemberMap.keySet()) {
					Member m = segmentToMemberMap.get(segmentNumber);
					LMMember lmMember = memberMap.get(m);
					SegmentMapping segmentMapping = SegmentMapping.newBuilder().setSegmentNumber(segmentNumber).setMember(lmMember).build();
					indexMappingBuilder.addSegmentMapping(segmentMapping);
				}
				responseBuilder.addIndexMapping(indexMappingBuilder);
			}

			return responseBuilder.build();
		}
		finally {
			globalLock.readLock().unlock();
		}
	}

}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy