CouchbaseWriterTest.java

/*
 * 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 gobblin.couchbase.writer;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentLinkedDeque;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.avro.Schema;
import org.apache.avro.SchemaBuilder;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericRecord;
import org.apache.commons.math3.util.Pair;
import org.apache.http.HttpResponse;
import org.apache.http.NameValuePair;
import org.apache.http.client.entity.UrlEncodedFormEntity;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.message.BasicNameValuePair;
import org.testng.Assert;
import org.testng.annotations.AfterSuite;
import org.testng.annotations.BeforeSuite;
import org.testng.annotations.Test;

import com.couchbase.client.deps.io.netty.buffer.ByteBuf;
import com.couchbase.client.java.Bucket;
import com.couchbase.client.java.document.AbstractDocument;
import com.couchbase.client.java.document.RawJsonDocument;
import com.couchbase.client.java.env.CouchbaseEnvironment;
import com.couchbase.client.java.env.DefaultCouchbaseEnvironment;
import com.google.gson.Gson;
import com.typesafe.config.Config;
import com.typesafe.config.ConfigFactory;

import lombok.extern.slf4j.Slf4j;

import gobblin.converter.Converter;
import gobblin.converter.DataConversionException;
import gobblin.couchbase.CouchbaseTestServer;
import gobblin.couchbase.common.TupleDocument;
import gobblin.couchbase.converter.AnyToCouchbaseJsonConverter;
import gobblin.couchbase.converter.AvroToCouchbaseTupleConverter;
import gobblin.metrics.RootMetricContext;
import gobblin.metrics.reporter.OutputStreamReporter;
import gobblin.test.TestUtils;
import gobblin.writer.AsyncWriterManager;
import gobblin.writer.WriteCallback;
import gobblin.writer.WriteResponse;

@Slf4j
public class CouchbaseWriterTest {

  private CouchbaseTestServer _couchbaseTestServer;
  private CouchbaseEnvironment _couchbaseEnvironment;

  @BeforeSuite
  public void startServers() {
    _couchbaseTestServer = new CouchbaseTestServer(TestUtils.findFreePort());
    _couchbaseTestServer.start();
    _couchbaseEnvironment = DefaultCouchbaseEnvironment.builder().bootstrapHttpEnabled(true)
        .bootstrapHttpDirectPort(_couchbaseTestServer.getPort())
        .bootstrapCarrierDirectPort(_couchbaseTestServer.getServerPort()).bootstrapCarrierEnabled(false)
        .kvTimeout(10000).build();
  }


  /**
   * Implement the equivalent of:
   * curl -XPOST -u Administrator:password localhost:httpPort/pools/default/buckets \ -d bucketType=couchbase \
   * -d name={@param bucketName} -d authType=sasl -d ramQuotaMB=200
   **/
  private boolean createBucket(String bucketName) {
    CloseableHttpClient httpClient = HttpClientBuilder.create().build();
    try {
      HttpPost httpPost = new HttpPost("http://localhost:" + _couchbaseTestServer.getPort() + "/pools/default/buckets");
      List<NameValuePair> params = new ArrayList<>(2);
      params.add(new BasicNameValuePair("bucketType", "couchbase"));
      params.add(new BasicNameValuePair("name", bucketName));
      params.add(new BasicNameValuePair("authType", "sasl"));
      params.add(new BasicNameValuePair("ramQuotaMB", "200"));
      httpPost.setEntity(new UrlEncodedFormEntity(params, "UTF-8"));
      //Execute and get the response.
      HttpResponse response = httpClient.execute(httpPost);
      log.info(String.valueOf(response.getStatusLine().getStatusCode()));
      return true;
    }
    catch (Exception e) {
      log.error("Failed to create bucket {}", bucketName, e);
      return false;
    }
  }

  @AfterSuite
  public void stopServers() {
    _couchbaseTestServer.stop();
  }

  /**
   * Test that a single tuple document can be written successfully.
   * @throws IOException
   * @throws DataConversionException
   * @throws ExecutionException
   * @throws InterruptedException
   */
  @Test
  public void testTupleDocumentWrite()
      throws IOException, DataConversionException, ExecutionException, InterruptedException {
    Properties props = new Properties();
    props.setProperty(CouchbaseWriterConfigurationKeys.BUCKET, "default");
    Config config = ConfigFactory.parseProperties(props);

    CouchbaseWriter writer = new CouchbaseWriter(_couchbaseEnvironment, config);
    try {
      Schema dataRecordSchema =
          SchemaBuilder.record("Data").fields().name("data").type().bytesType().noDefault().name("flags").type().intType()
              .noDefault().endRecord();

      Schema schema = SchemaBuilder.record("TestRecord").fields().name("key").type().stringType().noDefault().name("data")
          .type(dataRecordSchema).noDefault().endRecord();

      GenericData.Record testRecord = new GenericData.Record(schema);

      String testContent = "hello world";

      GenericData.Record dataRecord = new GenericData.Record(dataRecordSchema);
      dataRecord.put("data", ByteBuffer.wrap(testContent.getBytes(Charset.forName("UTF-8"))));
      dataRecord.put("flags", 0);

      testRecord.put("key", "hello");
      testRecord.put("data", dataRecord);

      Converter<Schema, String, GenericRecord, TupleDocument> recordConverter = new AvroToCouchbaseTupleConverter();

      TupleDocument doc = recordConverter.convertRecord("", testRecord, null).iterator().next();
      writer.write(doc, null).get();
      TupleDocument returnDoc = writer.getBucket().get("hello", TupleDocument.class);

      byte[] returnedBytes = new byte[returnDoc.content().value1().readableBytes()];
      returnDoc.content().value1().readBytes(returnedBytes);
      Assert.assertEquals(returnedBytes, testContent.getBytes(Charset.forName("UTF-8")));

      int returnedFlags = returnDoc.content().value2();
      Assert.assertEquals(returnedFlags, 0);

    } finally {
      writer.close();
    }

  }

  /**
   * Test that a single Json document can be written successfully
   * @throws IOException
   * @throws DataConversionException
   * @throws ExecutionException
   * @throws InterruptedException
   */
  @Test(groups={"timeout"})
  public void testJsonDocumentWrite()
      throws IOException, DataConversionException, ExecutionException, InterruptedException {
    CouchbaseWriter writer = new CouchbaseWriter(_couchbaseEnvironment, ConfigFactory.empty());
    try {

      String key = "hello";
      String testContent = "hello world";
      HashMap<String, String> contentMap = new HashMap<>();
      contentMap.put("value", testContent);
      Gson gson = new Gson();
      String jsonString = gson.toJson(contentMap);
      RawJsonDocument jsonDocument = RawJsonDocument.create(key, jsonString);
      writer.write(jsonDocument, null).get();
      RawJsonDocument returnDoc = writer.getBucket().get(key, RawJsonDocument.class);

      Map<String, String> returnedMap = gson.fromJson(returnDoc.content(), Map.class);
      Assert.assertEquals(testContent, returnedMap.get("value"));
    } finally {
      writer.close();
    }
  }

  private void drainQueue(BlockingQueue<Pair<AbstractDocument, Future>> queue, int threshold, long sleepTime,
      TimeUnit sleepUnit, List<Pair<AbstractDocument, Future>> failedFutures) {
    while (queue.remainingCapacity() < threshold) {
      if (sleepTime > 0) {
        Pair<AbstractDocument, Future> topElement = queue.peek();
        if (topElement != null) {
          try {
            topElement.getSecond().get(sleepTime, sleepUnit);
          } catch (Exception te) {
            failedFutures.add(topElement);
          }
          queue.poll();
        }
      }
    }
  }

  /**
   * An iterator that applies the {@link AnyToCouchbaseJsonConverter} converter to Objects
   */
  class JsonDocumentIterator implements Iterator<AbstractDocument> {
    private final int _maxRecords;
    private int _currRecord;
    private Iterator<Object> _objectIterator;
    private final Converter<String, String, Object, RawJsonDocument> _recordConverter =
        new AnyToCouchbaseJsonConverter();

    JsonDocumentIterator(Iterator<Object> genericRecordIterator) {
      this(genericRecordIterator, -1);
    }

    JsonDocumentIterator(Iterator<Object> genericRecordIterator, int maxRecords) {
      _objectIterator = genericRecordIterator;
      _maxRecords = maxRecords;
      _currRecord = 0;
    }

    @Override
    public boolean hasNext() {
      if (_maxRecords < 0) {
        return _objectIterator.hasNext();
      } else {
        return _objectIterator.hasNext() && (_currRecord < _maxRecords);
      }
    }

    @Override
    public AbstractDocument next() {
      _currRecord++;
      Object record = _objectIterator.next();
      try {
        return _recordConverter.convertRecord("", record, null).iterator().next();
      } catch (Exception e) {
        throw new RuntimeException(e);
      }
    }

    @Override
    public void remove() {
    }

  }

  /**
   * An iterator that applies the {@link AvroToCouchbaseTupleConverter} converter to GenericRecords
   */
  class TupleDocumentIterator implements Iterator<AbstractDocument> {
    private final int _maxRecords;
    private int _currRecord;
    private Iterator<GenericRecord> _genericRecordIterator;
    private final Converter<Schema, String, GenericRecord, TupleDocument> _recordConverter =
        new AvroToCouchbaseTupleConverter();

    TupleDocumentIterator(Iterator<GenericRecord> genericRecordIterator) {
      this(genericRecordIterator, -1);
    }

    TupleDocumentIterator(Iterator<GenericRecord> genericRecordIterator, int maxRecords) {
      _genericRecordIterator = genericRecordIterator;
      _maxRecords = maxRecords;
      _currRecord = 0;
    }

    @Override
    public boolean hasNext() {
      if (_maxRecords < 0) {
        return _genericRecordIterator.hasNext();
      } else {
        return _genericRecordIterator.hasNext() && (_currRecord < _maxRecords);
      }
    }

    @Override
    public TupleDocument next() {
      _currRecord++;
      GenericRecord record = _genericRecordIterator.next();
      try {
        return _recordConverter.convertRecord("", record, null).iterator().next();
      } catch (Exception e) {
        throw new RuntimeException(e);
      }
    }

    @Override
    public void remove() {
    }

  }

  class Verifier {
    private final Map<String, byte[]> verificationCache = new HashMap<>(1000);
    private Class recordClass;

    void onWrite(AbstractDocument doc)
        throws UnsupportedEncodingException {
      recordClass = doc.getClass();
      if (doc instanceof TupleDocument) {
        ByteBuf outgoingBuf = (((TupleDocument) doc).content()).value1();
        byte[] outgoingBytes = new byte[outgoingBuf.readableBytes()];
        outgoingBuf.getBytes(0, outgoingBytes);
        verificationCache.put(doc.id(), outgoingBytes);
      } else if (doc instanceof RawJsonDocument) {
        verificationCache.put(doc.id(), ((RawJsonDocument) doc).content().getBytes("UTF-8"));
      } else {
        throw new UnsupportedOperationException("Can only support TupleDocument or RawJsonDocument at this time");
      }
    }

    void verify(Bucket bucket)
        throws UnsupportedEncodingException {
      // verify
      System.out.println("Starting verification procedure");
      for (Map.Entry<String, byte[]> cacheEntry : verificationCache.entrySet()) {
        Object doc = bucket.get(cacheEntry.getKey(), recordClass);
        if (doc instanceof TupleDocument) {
          ByteBuf returnedBuf = (((TupleDocument) doc).content()).value1();
          byte[] returnedBytes = new byte[returnedBuf.readableBytes()];
          returnedBuf.getBytes(0, returnedBytes);
          Assert.assertEquals(returnedBytes, cacheEntry.getValue(), "Returned content for TupleDoc should be equal");
        } else if (doc instanceof RawJsonDocument) {
          byte[] returnedBytes = ((RawJsonDocument) doc).content().getBytes("UTF-8");
          Assert.assertEquals(returnedBytes, cacheEntry.getValue(), "Returned content for JsonDoc should be equal");
        } else {
          Assert.fail("Returned type was neither TupleDocument nor RawJsonDocument");
        }
      }
      System.out.println("Verification success!");
    }
  }

  /**
   * Uses the {@link AsyncWriterManager} to write records through a couchbase writer
   * It keeps a copy of the key, value combinations written and checks after all the writes have gone through.
   * @param recordIterator
   * @throws IOException
   */
  private void writeRecordsWithAsyncWriter(Iterator<AbstractDocument> recordIterator)
      throws IOException {
    boolean verbose = false;
    Properties props = new Properties();
    props.setProperty(CouchbaseWriterConfigurationKeys.BUCKET, "default");
    Config config = ConfigFactory.parseProperties(props);

    CouchbaseWriter writer = new CouchbaseWriter(_couchbaseEnvironment, config);

    try {
      AsyncWriterManager asyncWriterManager =
          AsyncWriterManager.builder().asyncDataWriter(writer).maxOutstandingWrites(100000).retriesEnabled(true)
              .numRetries(5).build();

      if (verbose) {
        // Create a reporter for metrics. This reporter will write metrics to STDOUT.
        OutputStreamReporter.Factory.newBuilder().build(new Properties());
        // Start all metric reporters.
        RootMetricContext.get().startReporting();
      }

      Verifier verifier = new Verifier();
      while (recordIterator.hasNext()) {
        AbstractDocument doc = recordIterator.next();
        verifier.onWrite(doc);
        asyncWriterManager.write(doc);
      }
      asyncWriterManager.commit();
      verifier.verify(writer.getBucket());
    } finally {
      writer.close();
    }
  }

  private List<Pair<AbstractDocument, Future>> writeRecords(Iterator<AbstractDocument> recordIterator,
      CouchbaseWriter writer, int outstandingRequests, long kvTimeout, TimeUnit kvTimeoutUnit)
      throws DataConversionException, UnsupportedEncodingException {
    final BlockingQueue<Pair<AbstractDocument, Future>> outstandingCallQueue =
        new LinkedBlockingDeque<>(outstandingRequests);
    final List<Pair<AbstractDocument, Future>> failedFutures = new ArrayList<>(outstandingRequests);

    int index = 0;
    long runTime = 0;
    final AtomicInteger callbackSuccesses = new AtomicInteger(0);
    final AtomicInteger callbackFailures = new AtomicInteger(0);
    final ConcurrentLinkedDeque<Throwable> callbackExceptions = new ConcurrentLinkedDeque<>();
    Verifier verifier = new Verifier();
    while (recordIterator.hasNext()) {
      AbstractDocument doc = recordIterator.next();
      index++;
      verifier.onWrite(doc);
      final long startTime = System.nanoTime();
      Future callFuture = writer.write(doc, new WriteCallback<TupleDocument>() {
        @Override
        public void onSuccess(WriteResponse<TupleDocument> writeResponse) {
          callbackSuccesses.incrementAndGet();
        }

        @Override
        public void onFailure(Throwable throwable) {
          callbackFailures.incrementAndGet();
          callbackExceptions.add(throwable);
        }
      });
      drainQueue(outstandingCallQueue, 1, kvTimeout, kvTimeoutUnit, failedFutures);
      outstandingCallQueue.add(new Pair<>(doc, callFuture));
      runTime += System.nanoTime() - startTime;
    }
    int failedWrites = 0;
    long responseStartTime = System.nanoTime();
    drainQueue(outstandingCallQueue, outstandingRequests, kvTimeout, kvTimeoutUnit, failedFutures);
    runTime += System.nanoTime() - responseStartTime;

    for (Throwable failure : callbackExceptions) {
      System.out.println(failure.getClass() + " : " + failure.getMessage());
    }
    failedWrites += failedFutures.size();

    System.out.println(
        "Total time to send " + index + " records = " + runTime / 1000000.0 + "ms, " + "Failed writes = " + failedWrites
            + " Callback Successes = " + callbackSuccesses.get() + "Callback Failures = " + callbackFailures.get());

    verifier.verify(writer.getBucket());
    return failedFutures;
  }

  @Test
  public void testMultiTupleDocumentWrite()
      throws IOException, DataConversionException, ExecutionException, InterruptedException {
    CouchbaseWriter writer = new CouchbaseWriter(_couchbaseEnvironment, ConfigFactory.empty());
    try {

      final Schema dataRecordSchema =
          SchemaBuilder.record("Data").fields().name("data").type().bytesType().noDefault().name("flags").type().intType()
              .noDefault().endRecord();

      final Schema schema =
          SchemaBuilder.record("TestRecord").fields().name("key").type().stringType().noDefault().name("data")
              .type(dataRecordSchema).noDefault().endRecord();

      final int numRecords = 1000;
      int outstandingRequests = 99;

      Iterator<GenericRecord> recordIterator = new Iterator<GenericRecord>() {
        private int currentIndex;

        @Override
        public void remove() {
        }

        @Override
        public boolean hasNext() {
          return (currentIndex < numRecords);
        }

        @Override
        public GenericRecord next() {
          GenericData.Record testRecord = new GenericData.Record(schema);

          String testContent = "hello world" + currentIndex;
          GenericData.Record dataRecord = new GenericData.Record(dataRecordSchema);
          dataRecord.put("data", ByteBuffer.wrap(testContent.getBytes(Charset.forName("UTF-8"))));
          dataRecord.put("flags", 0);

          testRecord.put("key", "hello" + currentIndex);
          testRecord.put("data", dataRecord);
          currentIndex++;
          return testRecord;
        }
      };

      long kvTimeout = 10000;
      TimeUnit kvTimeoutUnit = TimeUnit.MILLISECONDS;
      writeRecords(new TupleDocumentIterator(recordIterator), writer, outstandingRequests, kvTimeout, kvTimeoutUnit);
    } finally {
      writer.close();
    }
  }

  @Test
  public void testMultiJsonDocumentWriteWithAsyncWriter()
      throws IOException, DataConversionException, ExecutionException, InterruptedException {

    final int numRecords = 1000;

    Iterator<Object> recordIterator = new Iterator<Object>() {
      private int currentIndex;

      @Override
      public boolean hasNext() {
        return (currentIndex < numRecords);
      }

      @Override
      public Object next() {

        String testContent = "hello world" + currentIndex;
        String key = "hello" + currentIndex;
        HashMap<String, String> contentMap = new HashMap<>();
        contentMap.put("key", key);
        contentMap.put("value", testContent);
        currentIndex++;
        return contentMap;
      }

      @Override
      public void remove() {
      }

    };
    writeRecordsWithAsyncWriter(new JsonDocumentIterator(recordIterator));
  }

  @Test
  public void testMultiTupleDocumentWriteWithAsyncWriter()
      throws IOException, DataConversionException, ExecutionException, InterruptedException {
    final Schema dataRecordSchema =
        SchemaBuilder.record("Data").fields().name("data").type().bytesType().noDefault().name("flags").type().intType()
            .noDefault().endRecord();

    final Schema schema =
        SchemaBuilder.record("TestRecord").fields().name("key").type().stringType().noDefault().name("data")
            .type(dataRecordSchema).noDefault().endRecord();

    final int numRecords = 1000;

    Iterator<GenericRecord> recordIterator = new Iterator<GenericRecord>() {
      private int currentIndex;

      @Override
      public void remove() {
      }

      @Override
      public boolean hasNext() {
        return (currentIndex < numRecords);
      }

      @Override
      public GenericRecord next() {
        GenericData.Record testRecord = new GenericData.Record(schema);

        String testContent = "hello world" + currentIndex;
        GenericData.Record dataRecord = new GenericData.Record(dataRecordSchema);
        dataRecord.put("data", ByteBuffer.wrap(testContent.getBytes(Charset.forName("UTF-8"))));
        dataRecord.put("flags", 0);

        testRecord.put("key", "hello" + currentIndex);
        testRecord.put("data", dataRecord);
        currentIndex++;
        return testRecord;
      }
    };
    writeRecordsWithAsyncWriter(new TupleDocumentIterator(recordIterator));
  }
}