D-Net

package eu.dnetlib.functionality.index.solr.feed;

import java.io.IOException;

import java.net.MalformedURLException;

import java.net.URI;

import java.net.URL;

import java.util.Collection;

import java.util.Iterator;

import java.util.List;

import java.util.regex.Matcher;

import java.util.regex.Pattern;

import org.apache.commons.lang.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.apache.solr.client.solrj.SolrServer;

import org.apache.solr.client.solrj.SolrServerException;

import org.apache.solr.client.solrj.impl.CloudSolrServer;

import org.apache.solr.client.solrj.impl.HttpSolrServer;

import org.apache.solr.client.solrj.response.UpdateResponse;

import org.apache.solr.common.SolrInputDocument;

import com.google.common.base.Splitter;

import com.google.common.collect.Lists;

import com.google.common.hash.HashFunction;

import com.google.common.hash.Hashing;

import eu.dnetlib.miscutils.functional.xml.DnetXsltFunctions;

public class SolrServerPool {

	/**

	 * logger.

	 */

	private static final Log log = LogFactory.getLog(SolrServerPool.class); // NOPMD by marko on 11/24/08 5:02 PM

	/*

	 * We run into problems when using the ConcurrentUpdateSolrServer, so we're sticking to the HttpSolrServer.

	 */

	private final List<HttpSolrServer> updateServerPool = Lists.newArrayList();

	private CloudSolrServer cloudServer = null;

	private final HashFunction hash = Hashing.murmur3_32();

	public SolrServerPool(final String updateUrlLocal, final String updateUrlList, final String zkHost, final String collection, final boolean localFeeding) {

		for (URL url : parseUrlListPattern(updateUrlLocal, updateUrlList, localFeeding)) {

			String collectionUrl = url + "/" + collection;

			log.info("creating new HttpSolrServer: " + collectionUrl);

			System.out.println("creating new HttpSolrServer: " + collectionUrl);

			updateServerPool.add(new HttpSolrServer(collectionUrl));

		}

		if (!StringUtils.isBlank(zkHost)) {

			log.info("creating new CloudSolrServer: " + zkHost);

			System.out.println("creating new CloudSolrServer: " + zkHost);

			cloudServer = new CloudSolrServer(zkHost);

			cloudServer.setDefaultCollection(collection);

		}

	}

	public UpdateResponse add(final SolrInputDocument doc) throws SolrServerException, IOException {

		return updateServerPool.get(hashPick(doc)).add(doc);

	}

	public UpdateResponse addAll(final Iterator<SolrInputDocument> docs) throws SolrServerException, IOException {

		if (updateServerPool.size() == 1) return updateServerPool.get(0).add(docs);

		int i = Integer.parseInt(DnetXsltFunctions.randomInt(updateServerPool.size()));

		return updateServerPool.get(i).add(docs);

	}

	public UpdateResponse addAll(final Collection<SolrInputDocument> docs) throws SolrServerException, IOException {

		if (updateServerPool.size() == 1) return updateServerPool.get(0).add(docs);

		int i = Integer.parseInt(DnetXsltFunctions.randomInt(updateServerPool.size()));

		return updateServerPool.get(i).add(docs);

	}

	public void deleteByQuery(final String query) throws SolrServerException, IOException {

		if (cloudServer != null) {

			cloudServer.deleteByQuery(query);

		} else {

			updateServerPool.get(0).deleteByQuery(query);

		}

	}

	public void commitAll() throws SolrServerException, IOException {

		if (cloudServer != null) {

			cloudServer.commit();

		} else {

			updateServerPool.get(0).commit();

		}

	}

	public void shutdownAll() throws SolrServerException {

		if (cloudServer != null) {

			cloudServer.shutdown();

		}

		for (SolrServer server : updateServerPool) {

			server.shutdown();

		}

	}

	// //////////////////

	private int hashPick(final SolrInputDocument doc) {

		final int hashCode = hash.hashBytes(doc.getFieldValue("__indexrecordidentifier").toString().getBytes()).asInt();

		return Math.abs(hashCode) % updateServerPool.size();

	}

	public List<URL> parseUrlListPattern(final String local, final String list, final boolean localFeeding) {

		final List<URL> res = Lists.newArrayList();

		try {

			if (localFeeding) {

				res.add(new URL(local));

			} else {

				Matcher matcher = Pattern.compile("(^.*)\\[(\\d+)\\.\\.(\\d+)\\](.*$)").matcher(list);

				if (matcher.matches()) {

					final String prefix = matcher.group(1);

					int lb = Integer.parseInt(matcher.group(2));

					int ub = Integer.parseInt(matcher.group(3));

					final String suffix = matcher.group(4);

					for (int i = lb; i <= ub; i++) {

						res.add(new URL(prefix + i + suffix));

					}

				}

			}

		} catch (MalformedURLException e) {

			throw new IllegalArgumentException("invalid url list: " + list, e);

		}

		log.info("parsed url(s): " + res);

		return res;

	}

	public List<URL> parseUrlList(final String list) throws MalformedURLException {

		final List<URL> res = Lists.newArrayList();

		for (final String url : Splitter.on(",").trimResults().split(list)) {

			res.add(URI.create(url).toURL());

		}

		return res;

	}

}