/*
 * 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.lucene.replicator;

import java.io.Closeable;
import java.io.IOException;
import java.util.HashMap;
import java.util.concurrent.Callable;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.SnapshotDeletionPolicy;
import org.apache.lucene.replicator.ReplicationClient.ReplicationHandler;
import org.apache.lucene.replicator.ReplicationClient.SourceDirectoryFactory;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.store.MockDirectoryWrapper;
import org.apache.lucene.tests.util.TestUtil;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.ThreadInterruptedException;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

public class TestIndexReplicationClient extends ReplicatorTestCase {

  private static class IndexReadyCallback implements Callable<Boolean>, Closeable {

    private final Directory indexDir;
    private DirectoryReader reader;
    private long lastGeneration = -1;

    public IndexReadyCallback(Directory indexDir) throws IOException {
      this.indexDir = indexDir;
      if (DirectoryReader.indexExists(indexDir)) {
        reader = DirectoryReader.open(indexDir);
        lastGeneration = reader.getIndexCommit().getGeneration();
      }
    }

    @Override
    public Boolean call() throws Exception {
      if (reader == null) {
        reader = DirectoryReader.open(indexDir);
        lastGeneration = reader.getIndexCommit().getGeneration();
      } else {
        DirectoryReader newReader = DirectoryReader.openIfChanged(reader);
        assertNotNull(
            "should not have reached here if no changes were made to the index", newReader);
        long newGeneration = newReader.getIndexCommit().getGeneration();
        assertTrue(
            "expected newer generation; current=" + lastGeneration + " new=" + newGeneration,
            newGeneration > lastGeneration);
        reader.close();
        reader = newReader;
        lastGeneration = newGeneration;
        TestUtil.checkIndex(indexDir);
      }
      return null;
    }

    @Override
    public void close() throws IOException {
      IOUtils.close(reader);
    }
  }

  private MockDirectoryWrapper publishDir, handlerDir;
  private Replicator replicator;
  private SourceDirectoryFactory sourceDirFactory;
  private ReplicationClient client;
  private ReplicationHandler handler;
  private IndexWriter publishWriter;
  private IndexReadyCallback callback;

  private static final String VERSION_ID = "version";

  private void assertHandlerRevision(int expectedID, Directory dir) throws IOException {
    // loop as long as client is alive. test-framework will terminate us if
    // there's a serious bug, e.g. client doesn't really update. otherwise,
    // introducing timeouts is not good, can easily lead to false positives.
    while (client.isUpdateThreadAlive()) {
      // give client a chance to update
      try {
        Thread.sleep(100);
      } catch (InterruptedException e) {
        throw new ThreadInterruptedException(e);
      }

      try {
        DirectoryReader reader = DirectoryReader.open(dir);
        try {
          int handlerID =
              Integer.parseInt(reader.getIndexCommit().getUserData().get(VERSION_ID), 16);
          if (expectedID == handlerID) {
            return;
          } else if (VERBOSE) {
            System.out.println(
                "expectedID="
                    + expectedID
                    + " actual="
                    + handlerID
                    + " generation="
                    + reader.getIndexCommit().getGeneration());
          }
        } finally {
          reader.close();
        }
      } catch (
          @SuppressWarnings("unused")
          Exception e) {
        // we can hit IndexNotFoundException or e.g. EOFException (on
        // segments_N) because it is being copied at the same time it is read by
        // DirectoryReader.open().
      }
    }
  }

  private Revision createRevision(final int id) throws IOException {
    publishWriter.addDocument(new Document());
    publishWriter.setLiveCommitData(
        new HashMap<String, String>() {
          {
            put(VERSION_ID, Integer.toString(id, 16));
          }
        }.entrySet());
    publishWriter.commit();
    return new IndexRevision(publishWriter);
  }

  @Override
  @Before
  public void setUp() throws Exception {
    super.setUp();
    publishDir = newMockDirectory();
    handlerDir = newMockDirectory();
    sourceDirFactory = new PerSessionDirectoryFactory(createTempDir("replicationClientTest"));
    replicator = new LocalReplicator();
    callback = new IndexReadyCallback(handlerDir);
    handler = new IndexReplicationHandler(handlerDir, callback);
    client = new ReplicationClient(replicator, handler, sourceDirFactory);

    IndexWriterConfig conf = newIndexWriterConfig(null);
    conf.setIndexDeletionPolicy(new SnapshotDeletionPolicy(conf.getIndexDeletionPolicy()));
    publishWriter = new IndexWriter(publishDir, conf);
  }

  @After
  @Override
  public void tearDown() throws Exception {
    publishWriter.close();
    IOUtils.close(client, callback, replicator, publishDir, handlerDir);
    super.tearDown();
  }

  @Test
  public void testNoUpdateThread() throws Exception {
    assertNull("no version expected at start", handler.currentVersion());

    // Callback validates the replicated index
    replicator.publish(createRevision(1));
    client.updateNow();

    replicator.publish(createRevision(2));
    client.updateNow();

    // Publish two revisions without update, handler should be upgraded to latest
    replicator.publish(createRevision(3));
    replicator.publish(createRevision(4));
    client.updateNow();
  }

  @Test
  public void testUpdateThread() throws Exception {
    client.startUpdateThread(10, "index");

    replicator.publish(createRevision(1));
    assertHandlerRevision(1, handlerDir);

    replicator.publish(createRevision(2));
    assertHandlerRevision(2, handlerDir);

    // Publish two revisions without update, handler should be upgraded to latest
    replicator.publish(createRevision(3));
    replicator.publish(createRevision(4));
    assertHandlerRevision(4, handlerDir);
  }

  @Test
  public void testRestart() throws Exception {
    replicator.publish(createRevision(1));
    client.updateNow();

    replicator.publish(createRevision(2));
    client.updateNow();

    client.stopUpdateThread();
    client.close();
    client = new ReplicationClient(replicator, handler, sourceDirFactory);

    // Publish two revisions without update, handler should be upgraded to latest
    replicator.publish(createRevision(3));
    replicator.publish(createRevision(4));
    client.updateNow();
  }

  /*
   * This test verifies that the client and handler do not end up in a corrupt
   * index if exceptions are thrown at any point during replication. Either when
   * a client copies files from the server to the temporary space, or when the
   * handler copies them to the index directory.
   */
  @Test
  public void testConsistencyOnExceptions() throws Exception {
    // so the handler's index isn't empty
    replicator.publish(createRevision(1));
    client.updateNow();
    client.close();
    callback.close();

    // wrap sourceDirFactory to return a MockDirWrapper so we can simulate errors
    final SourceDirectoryFactory in = sourceDirFactory;
    final AtomicInteger failures = new AtomicInteger(atLeast(10));
    sourceDirFactory =
        new SourceDirectoryFactory() {

          private long clientMaxSize = 100, handlerMaxSize = 100;
          private double clientExRate = 1.0, handlerExRate = 1.0;

          @Override
          public void cleanupSession(String sessionID) throws IOException {
            in.cleanupSession(sessionID);
          }

          @SuppressWarnings("synthetic-access")
          @Override
          public Directory getDirectory(String sessionID, String source) throws IOException {
            Directory dir = in.getDirectory(sessionID, source);
            if (random().nextBoolean()
                && failures.get() > 0) { // client should fail, return wrapped dir
              MockDirectoryWrapper mdw = new MockDirectoryWrapper(random(), dir);
              mdw.setRandomIOExceptionRateOnOpen(clientExRate);
              mdw.setMaxSizeInBytes(clientMaxSize);
              mdw.setRandomIOExceptionRate(clientExRate);
              mdw.setCheckIndexOnClose(false);
              clientMaxSize *= 2;
              clientExRate /= 2;
              return mdw;
            }

            if (failures.get() > 0 && random().nextBoolean()) { // handler should fail
              handlerDir.setMaxSizeInBytes(handlerMaxSize);
              handlerDir.setRandomIOExceptionRateOnOpen(handlerExRate);
              handlerDir.setRandomIOExceptionRate(handlerExRate);
              handlerMaxSize *= 2;
              handlerExRate /= 2;
            } else {
              // disable errors
              handlerDir.setMaxSizeInBytes(0);
              handlerDir.setRandomIOExceptionRate(0.0);
              handlerDir.setRandomIOExceptionRateOnOpen(0.0);
            }
            return dir;
          }
        };

    handler =
        new IndexReplicationHandler(
            handlerDir,
            new Callable<Boolean>() {
              @Override
              public Boolean call() throws Exception {
                if (random().nextDouble() < 0.2 && failures.get() > 0) {
                  throw new RuntimeException("random exception from callback");
                }
                return null;
              }
            });

    // wrap handleUpdateException so we can act on the thrown exception
    client =
        new ReplicationClient(replicator, handler, sourceDirFactory) {
          @SuppressWarnings("synthetic-access")
          @Override
          protected void handleUpdateException(Throwable t) {
            if (t instanceof IOException) {
              if (VERBOSE) {
                System.out.println("hit exception during update: " + t);
                t.printStackTrace(System.out);
              }
              try {
                // test that the index can be read and also some basic statistics
                DirectoryReader reader = DirectoryReader.open(handlerDir.getDelegate());
                try {
                  int numDocs = reader.numDocs();
                  int version =
                      Integer.parseInt(reader.getIndexCommit().getUserData().get(VERSION_ID), 16);
                  assertEquals(numDocs, version);
                } finally {
                  reader.close();
                }
                // verify index consistency
                TestUtil.checkIndex(handlerDir.getDelegate());
              } catch (IOException e) {
                // exceptions here are bad, don't ignore them
                throw new RuntimeException(e);
              } finally {
                // count-down number of failures
                failures.decrementAndGet();
                assert failures.get() >= 0 : "handler failed too many times: " + failures.get();
                if (VERBOSE) {
                  if (failures.get() == 0) {
                    System.out.println("no more failures expected");
                  } else {
                    System.out.println("num failures left: " + failures.get());
                  }
                }
              }
            } else {
              if (t instanceof RuntimeException) throw (RuntimeException) t;
              throw new RuntimeException(t);
            }
          }
        };

    client.startUpdateThread(10, "index");

    final Directory baseHandlerDir = handlerDir.getDelegate();
    int numRevisions = atLeast(20);
    for (int i = 2; i < numRevisions; i++) {
      replicator.publish(createRevision(i));
      assertHandlerRevision(i, baseHandlerDir);
    }

    // disable errors -- maybe randomness didn't exhaust all allowed failures,
    // and we don't want e.g. CheckIndex to hit false errors.
    handlerDir.setMaxSizeInBytes(0);
    handlerDir.setRandomIOExceptionRate(0.0);
    handlerDir.setRandomIOExceptionRateOnOpen(0.0);
  }
}
