org.apache.solr.client.solrj.impl.StreamingUpdateSolrServer.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.solr.client.solrj.impl.StreamingUpdateSolrServer.java

Source

/**
 * 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 org.apache.solr.client.solrj.impl;

import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.net.MalformedURLException;
import java.util.LinkedList;
import java.util.Queue;
import java.util.concurrent.*;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.apache.commons.httpclient.HttpClient;
import org.apache.commons.httpclient.HttpStatus;
import org.apache.commons.httpclient.methods.PostMethod;
import org.apache.commons.httpclient.methods.RequestEntity;
import org.apache.solr.client.solrj.SolrRequest;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.client.solrj.request.UpdateRequest;
import org.apache.solr.client.solrj.util.ClientUtils;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.params.UpdateParams;
import org.apache.solr.common.util.NamedList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * {@link StreamingUpdateSolrServer} buffers all added documents and writes them
 * into open HTTP connections. This class is thread safe.
 * 
 * Although any SolrServer request can be made with this implementation, 
 * it is only recommended to use the {@link StreamingUpdateSolrServer} with
 * /update requests.  The query interface is better suited for 
 * 
 * @version $Id: CommonsHttpSolrServer.java 724175 2008-12-07 19:07:11Z ryan $
 * @since solr 1.4
 */
public class StreamingUpdateSolrServer extends CommonsHttpSolrServer {
    static final Logger log = LoggerFactory.getLogger(StreamingUpdateSolrServer.class);

    final BlockingQueue<UpdateRequest> queue;
    final ExecutorService scheduler = Executors.newCachedThreadPool();
    final String updateUrl = "/update";
    final Queue<Runner> runners;
    volatile CountDownLatch lock = null; // used to block everything
    final int threadCount;

    /**
     * Uses an internal MultiThreadedHttpConnectionManager to manage http connections
     *
     * @param solrServerUrl The Solr server URL
     * @param queueSize     The buffer size before the documents are sent to the server
     * @param threadCount   The number of background threads used to empty the queue
     * @throws MalformedURLException
     */
    public StreamingUpdateSolrServer(String solrServerUrl, int queueSize, int threadCount)
            throws MalformedURLException {
        this(solrServerUrl, null, queueSize, threadCount);
    }

    /**
     * Uses the supplied HttpClient to send documents to the Solr server, the HttpClient should be instantiated using a
     * MultiThreadedHttpConnectionManager.
     */
    public StreamingUpdateSolrServer(String solrServerUrl, HttpClient client, int queueSize, int threadCount)
            throws MalformedURLException {
        super(solrServerUrl, client);
        queue = new LinkedBlockingQueue<UpdateRequest>(queueSize);
        this.threadCount = threadCount;
        runners = new LinkedList<Runner>();
    }

    /**
     * Opens a connection and sends everything...
     */
    class Runner implements Runnable {
        final Lock runnerLock = new ReentrantLock();

        public void run() {
            runnerLock.lock();

            // info is ok since this should only happen once for each thread
            log.info("starting runner: {}", this);
            PostMethod method = null;
            try {
                do {
                    try {
                        RequestEntity request = new RequestEntity() {
                            // we don't know the length
                            public long getContentLength() {
                                return -1;
                            }

                            public String getContentType() {
                                return ClientUtils.TEXT_XML;
                            }

                            public boolean isRepeatable() {
                                return false;
                            }

                            public void writeRequest(OutputStream out) throws IOException {
                                try {
                                    OutputStreamWriter writer = new OutputStreamWriter(out, "UTF-8");
                                    writer.append("<stream>"); // can be anything...
                                    UpdateRequest req = queue.poll(250, TimeUnit.MILLISECONDS);
                                    while (req != null) {
                                        log.debug("sending: {}", req);
                                        req.writeXML(writer);

                                        // check for commit or optimize
                                        SolrParams params = req.getParams();
                                        if (params != null) {
                                            String fmt = null;
                                            if (params.getBool(UpdateParams.OPTIMIZE, false)) {
                                                fmt = "<optimize waitSearcher=\"%s\" waitFlush=\"%s\" />";
                                            } else if (params.getBool(UpdateParams.COMMIT, false)) {
                                                fmt = "<commit waitSearcher=\"%s\" waitFlush=\"%s\" />";
                                            }
                                            if (fmt != null) {
                                                log.info(fmt);
                                                writer.write(String.format(fmt,
                                                        params.getBool(UpdateParams.WAIT_SEARCHER, false) + "",
                                                        params.getBool(UpdateParams.WAIT_FLUSH, false) + ""));
                                            }
                                        }

                                        writer.flush();
                                        req = queue.poll(250, TimeUnit.MILLISECONDS);
                                    }
                                    writer.append("</stream>");
                                    writer.flush();
                                } catch (InterruptedException e) {
                                    e.printStackTrace();
                                }
                            }
                        };

                        method = new PostMethod(_baseURL + updateUrl);
                        method.setRequestEntity(request);
                        method.setFollowRedirects(false);
                        method.addRequestHeader("User-Agent", AGENT);

                        int statusCode = getHttpClient().executeMethod(method);
                        if (statusCode != HttpStatus.SC_OK) {
                            StringBuilder msg = new StringBuilder();
                            msg.append(method.getStatusLine().getReasonPhrase());
                            msg.append("\n\n");
                            msg.append(method.getStatusText());
                            msg.append("\n\n");
                            msg.append("request: " + method.getURI());
                            handleError(new Exception(msg.toString()));
                        }
                    } finally {
                        try {
                            // make sure to release the connection
                            if (method != null)
                                method.releaseConnection();
                        } catch (Exception ex) {
                        }
                    }
                } while (!queue.isEmpty());
            } catch (Throwable e) {
                handleError(e);
            } finally {

                // remove it from the list of running things unless we are the last runner and the queue is full...
                // in which case, the next queue.put() would block and there would be no runners to handle it.
                // This case has been further handled by using offer instead of put, and using a retry loop
                // to avoid blocking forever (see request()).
                synchronized (runners) {
                    if (runners.size() == 1 && queue.remainingCapacity() == 0) {
                        // keep this runner alive
                        scheduler.execute(this);
                    } else {
                        runners.remove(this);
                    }
                }

                log.info("finished: {}", this);
                runnerLock.unlock();
            }
        }
    }

    @Override
    public NamedList<Object> request(final SolrRequest request) throws SolrServerException, IOException {
        if (!(request instanceof UpdateRequest)) {
            return super.request(request);
        }
        UpdateRequest req = (UpdateRequest) request;

        // this happens for commit...
        if (req.getDocuments() == null || req.getDocuments().isEmpty()) {
            blockUntilFinished();
            return super.request(request);
        }

        SolrParams params = req.getParams();
        if (params != null) {
            // check if it is waiting for the searcher
            if (params.getBool(UpdateParams.WAIT_SEARCHER, false)) {
                log.info("blocking for commit/optimize");
                blockUntilFinished(); // empty the queue
                return super.request(request);
            }
        }

        try {
            CountDownLatch tmpLock = lock;
            if (tmpLock != null) {
                tmpLock.await();
            }

            boolean success = queue.offer(req);

            for (;;) {
                synchronized (runners) {
                    if (runners.isEmpty() || (queue.remainingCapacity() < queue.size() // queue is half full and we can add more runners
                            && runners.size() < threadCount)) {
                        // We need more runners, so start a new one.
                        Runner r = new Runner();
                        runners.add(r);
                        scheduler.execute(r);
                    } else {
                        // break out of the retry loop if we added the element to the queue successfully, *and*
                        // while we are still holding the runners lock to prevent race conditions.
                        // race conditions.
                        if (success)
                            break;
                    }
                }

                // Retry to add to the queue w/o the runners lock held (else we risk temporary deadlock)
                // This retry could also fail because
                // 1) existing runners were not able to take off any new elements in the queue
                // 2) the queue was filled back up since our last try
                // If we succeed, the queue may have been completely emptied, and all runners stopped.
                // In all cases, we should loop back to the top to see if we need to start more runners.
                //
                if (!success) {
                    success = queue.offer(req, 100, TimeUnit.MILLISECONDS);
                }

            }

        } catch (InterruptedException e) {
            log.error("interrupted", e);
            throw new IOException(e.getLocalizedMessage());
        }

        // RETURN A DUMMY result
        NamedList<Object> dummy = new NamedList<Object>();
        dummy.add("NOTE", "the request is processed in a background stream");
        return dummy;
    }

    public synchronized void blockUntilFinished() {
        lock = new CountDownLatch(1);
        try {
            // Wait until no runners are running
            for (;;) {
                Runner runner;
                synchronized (runners) {
                    runner = runners.peek();
                }
                if (runner == null)
                    break;
                runner.runnerLock.lock();
                runner.runnerLock.unlock();
            }
        } finally {
            lock.countDown();
            lock = null;
        }
    }

    public void handleError(Throwable ex) {
        log.error("error", ex);
    }
}