More changes based on PR comments.

Author: tdas

streaming/src/main/java/org/apache/spark/streaming/util/WriteAheadLog.java

@@ -21,39 +21,39 @@
 import java.util.Iterator;
 
 /**
- * Interface representing a write ahead log (aka journal) that is used by Spark Streaming to
- * save the received data (by receivers) and associated metadata to a reliable storage, so that
- * they can be recovered after driver failures. See the Spark docs for more information on how
- * to plug in your own custom implementation of a write ahead log.
+ * This abstract class represents a write ahead log (aka journal) that is used by Spark Streaming
+ * to save the received data (by receivers) and associated metadata to a reliable storage, so that
+ * they can be recovered after driver failures. See the Spark documentation for more information
+ * on how to plug in your own custom implementation of a write ahead log.
  */
 @org.apache.spark.annotation.DeveloperApi
-public interface WriteAheadLog {
+public abstract class WriteAheadLog {
   /**
    * Write the record to the log and return the segment information that is necessary to read
    * back the written record. The time is used to the index the record, such that it can be
    * cleaned later. Note that the written data must be durable and readable (using the
    * segment info) by the time this function returns.
    */
-  WriteAheadLogSegment write(ByteBuffer record, long time);
+  abstract public WriteAheadLogRecordHandle write(ByteBuffer record, long time);
 
   /**
    * Read a written record based on the given segment information.
    */
-  ByteBuffer read(WriteAheadLogSegment segment);
+  abstract public ByteBuffer read(WriteAheadLogRecordHandle handle);
 
   /**
-   * Read and return an iterator of all the records that have written and not yet cleanup.
+   * Read and return an iterator of all the records that have been written but not yet cleaned up.
    */
-  Iterator<ByteBuffer> readAll();
+  abstract public Iterator<ByteBuffer> readAll();
 
   /**
-   * Cleanup all the records that are older than the given threshold time. It can wait for
+   * Clean all the records that are older than the threshold time. It can wait for
    * the completion of the deletion.
    */
-  void cleanup(long threshTime, boolean waitForCompletion);
+  abstract public void clean(long threshTime, boolean waitForCompletion);
 
   /**
    * Close this log and release any resources.
    */
-  void close();
+  abstract public void close();
 }

…streaming/util/WriteAheadLogSegment.java …ming/util/WriteAheadLogRecordHandle.javastreaming/src/main/java/org/apache/spark/streaming/util/WriteAheadLogSegment.java renamed to streaming/src/main/java/org/apache/spark/streaming/util/WriteAheadLogRecordHandle.java streaming/src/main/java/org/apache/spark/streaming/util/WriteAheadLogSegment.java renamed to streaming/src/main/java/org/apache/spark/streaming/util/WriteAheadLogRecordHandle.java

@@ -18,9 +18,13 @@
 package org.apache.spark.streaming.util;
 
 /**
- * This is an interface that represent the information required by any implementation of
- * a WriteAheadLog to read a written record.
+ * This abstract class represents a handle that refers to a record written in a
+ * {@link org.apache.spark.streaming.util.WriteAheadLog WriteAheadLog}.
+ * It must contain all the information necessary for the record to be read and returned by
+ * an implemenation of the WriteAheadLog class.
+ *
+ * @see org.apache.spark.streaming.util.WriteAheadLog
  */
 @org.apache.spark.annotation.DeveloperApi
-public interface WriteAheadLogSegment extends java.io.Serializable {
+public abstract class WriteAheadLogRecordHandle implements java.io.Serializable {
 }

streaming/src/main/scala/org/apache/spark/streaming/rdd/WriteAheadLogBackedBlockRDD.scala

@@ -40,7 +40,7 @@ private[streaming]
 class WriteAheadLogBackedBlockRDDPartition(
     val index: Int,
     val blockId: BlockId,
-    val segment: WriteAheadLogSegment)
+    val segment: WriteAheadLogRecordHandle)
   extends Partition
 
 
@@ -61,7 +61,7 @@ private[streaming]
 class WriteAheadLogBackedBlockRDD[T: ClassTag](
     @transient sc: SparkContext,
     @transient blockIds: Array[BlockId],
-    @transient segments: Array[WriteAheadLogSegment],
+    @transient segments: Array[WriteAheadLogRecordHandle],
     storeInBlockManager: Boolean,
     storageLevel: StorageLevel)
   extends BlockRDD[T](sc, blockIds) {

streaming/src/main/scala/org/apache/spark/streaming/receiver/ReceivedBlockHandler.scala

@@ -26,7 +26,7 @@ import org.apache.hadoop.fs.Path
 
 import org.apache.spark.storage._
 import org.apache.spark.streaming.receiver.WriteAheadLogBasedBlockHandler._
-import org.apache.spark.streaming.util.{WriteAheadLogSegment, WriteAheadLogUtils}
+import org.apache.spark.streaming.util.{WriteAheadLogRecordHandle, WriteAheadLogUtils}
 import org.apache.spark.util.{Clock, SystemClock, ThreadUtils}
 import org.apache.spark.{Logging, SparkConf, SparkException}
 
@@ -96,7 +96,7 @@ private[streaming] class BlockManagerBasedBlockHandler(
  */
 private[streaming] case class WriteAheadLogBasedStoreResult(
     blockId: StreamBlockId,
-    segment: WriteAheadLogSegment
+    segment: WriteAheadLogRecordHandle
   ) extends ReceivedBlockStoreResult
 
 
@@ -185,7 +185,7 @@ private[streaming] class WriteAheadLogBasedBlockHandler(
   }
 
   def cleanupOldBlocks(threshTime: Long) {
-    writeAheadLog.cleanup(threshTime, false)
+    writeAheadLog.clean(threshTime, false)
   }
 
   def stop() {

streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceivedBlockTracker.scala

@@ -155,7 +155,7 @@ private[streaming] class ReceivedBlockTracker(
     logInfo("Deleting batches " + timesToCleanup)
     writeToLog(BatchCleanupEvent(timesToCleanup))
     timeToAllocatedBlocks --= timesToCleanup
-    writeAheadLogOption.foreach(_.cleanup(cleanupThreshTime.milliseconds, waitForCompletion))
+    writeAheadLogOption.foreach(_.clean(cleanupThreshTime.milliseconds, waitForCompletion))
   }
 
   /** Stop the block tracker. */

streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLog.scala

@@ -95,7 +95,7 @@ private[streaming] class FileBasedWriteAheadLog(
     fileSegment
   }
 
-  def read(segment: WriteAheadLogSegment): ByteBuffer = {
+  def read(segment: WriteAheadLogRecordHandle): ByteBuffer = {
     val fileSegment = segment.asInstanceOf[FileBasedWriteAheadLogSegment]
     var reader: FileBasedWriteAheadLogRandomReader = null
     var byteBuffer: ByteBuffer = null
@@ -140,7 +140,7 @@ private[streaming] class FileBasedWriteAheadLog(
    * deleted. This should be set to true only for testing. Else the files will be deleted
    * asynchronously.
    */
-  def cleanup(threshTime: Long, waitForCompletion: Boolean): Unit = {
+  def clean(threshTime: Long, waitForCompletion: Boolean): Unit = {
     val oldLogFiles = synchronized { pastLogs.filter { _.endTime < threshTime } }
     logInfo(s"Attempting to clear ${oldLogFiles.size} old log files in $logDirectory " +
       s"older than $threshTime: ${oldLogFiles.map { _.path }.mkString("\n")}")

streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLogSegment.scala

@@ -18,4 +18,4 @@ package org.apache.spark.streaming.util
 
 /** Class for representing a segment of data in a write ahead log file */
 private[streaming] case class FileBasedWriteAheadLogSegment(path: String, offset: Long, length: Int)
-  extends WriteAheadLogSegment
+  extends WriteAheadLogRecordHandle

streaming/src/test/java/org/apache/spark/streaming/JavaWriteAheadLogSuite.java

@@ -0,0 +1,129 @@
+/*
+ * 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.spark.streaming;
+
+import java.util.ArrayList;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collection;
+
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.collections.Transformer;
+import org.apache.spark.SparkConf;
+import org.apache.spark.streaming.util.WriteAheadLog;
+import org.apache.spark.streaming.util.WriteAheadLogRecordHandle;
+import org.apache.spark.streaming.util.WriteAheadLogUtils;
+
+import org.junit.Test;
+import org.junit.Assert;
+
+class JavaWriteAheadLogSuiteHandle extends WriteAheadLogRecordHandle {
+  int index = -1;
+  public JavaWriteAheadLogSuiteHandle(int idx) {
+    index = idx;
+  }
+}
+
+public class JavaWriteAheadLogSuite extends WriteAheadLog {
+
+  class Record {
+    long time;
+    int index;
+    ByteBuffer buffer;
+
+    public Record(long tym, int idx, ByteBuffer buf) {
+      index = idx;
+      time = tym;
+      buffer = buf;
+    }
+  }
+  private int index = -1;
+  private ArrayList<Record> records = new ArrayList<Record>();
+
+
+  // Methods for WriteAheadLog
+  @Override
+  public WriteAheadLogRecordHandle write(java.nio.ByteBuffer record, long time) {
+    index += 1;
+    records.add(new org.apache.spark.streaming.JavaWriteAheadLogSuite.Record(time, index, record));
+    return new JavaWriteAheadLogSuiteHandle(index);
+  }
+
+  @Override
+  public java.nio.ByteBuffer read(WriteAheadLogRecordHandle handle) {
+    if (handle instanceof JavaWriteAheadLogSuiteHandle) {
+      int reqdIndex = ((JavaWriteAheadLogSuiteHandle) handle).index;
+      for (Record record: records) {
+        if (record.index == reqdIndex) {
+          return record.buffer;
+        }
+      }
+    }
+    return null;
+  }
+
+  @Override
+  public java.util.Iterator<java.nio.ByteBuffer> readAll() {
+    Collection<ByteBuffer> buffers = CollectionUtils.collect(records, new Transformer() {
+      @Override
+      public Object transform(Object input) {
+        return ((Record) input).buffer;
+      }
+    });
+    return buffers.iterator();
+  }
+
+  @Override
+  public void clean(long threshTime, boolean waitForCompletion) {
+    for (int i = 0; i < records.size(); i++) {
+      if (records.get(i).time < threshTime) {
+        records.remove(i);
+        i--;
+      }
+    }
+  }
+
+  @Override
+  public void close() {
+    records.clear();
+  }
+
+  @Test
+  public void testCustomWAL() {
+    SparkConf conf = new SparkConf();
+    conf.set("spark.streaming.driver.writeAheadLog.class", JavaWriteAheadLogSuite.class.getName());
+    WriteAheadLog wal = WriteAheadLogUtils.createLogForDriver(conf, null, null);
+
+    String data1 = "data1";
+    WriteAheadLogRecordHandle handle = wal.write(ByteBuffer.wrap(data1.getBytes()), 1234);
+    Assert.assertTrue(handle instanceof JavaWriteAheadLogSuiteHandle);
+    Assert.assertTrue(new String(wal.read(handle).array()).equals(data1));
+
+    wal.write(ByteBuffer.wrap("data2".getBytes()), 1235);
+    wal.write(ByteBuffer.wrap("data3".getBytes()), 1236);
+    wal.write(ByteBuffer.wrap("data4".getBytes()), 1237);
+    wal.clean(1236, false);
+
+    java.util.Iterator<java.nio.ByteBuffer> dataIterator = wal.readAll();
+    ArrayList<String> readData = new ArrayList<String>();
+    while (dataIterator.hasNext()) {
+      readData.add(new String(dataIterator.next().array()));
+    }
+    Assert.assertTrue(readData.equals(Arrays.asList("data3", "data4")));
+  }
+}

streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockHandlerSuite.scala

@@ -149,7 +149,7 @@ class ReceivedBlockHandlerSuite extends FunSuite with BeforeAndAfter with Matche
     }
   }
 
-  test("WriteAheadLogBasedBlockHandler - cleanup old blocks") {
+  test("WriteAheadLogBasedBlockHandler - clean old blocks") {
     withWriteAheadLogBasedBlockHandler { handler =>
       val blocks = Seq.tabulate(10) { i => IteratorBlock(Iterator(1 to i)) }
       storeBlocks(handler, blocks)

streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala

@@ -88,7 +88,7 @@ class ReceivedBlockTrackerSuite
     receivedBlockTracker.getUnallocatedBlocks(streamId) shouldEqual blockInfos
   }
 
-  test("block addition, block to batch allocation and cleanup with write ahead log") {
+  test("block addition, block to batch allocation and clean up with write ahead log") {
     val manualClock = new ManualClock
     // Set the time increment level to twice the rotation interval so that every increment creates
     // a new log file
@@ -175,7 +175,7 @@ class ReceivedBlockTrackerSuite
     eventually(timeout(10 seconds), interval(10 millisecond)) {
       getWriteAheadLogFiles() should not contain oldestLogFile
     }
-    printLogFiles("After cleanup")
+    printLogFiles("After clean")
 
     // Restart tracker and verify recovered state, specifically whether info about the first
     // batch has been removed, but not the second batch
@@ -206,7 +206,7 @@ class ReceivedBlockTrackerSuite
 
   /**
    * Create tracker object with the optional provided clock. Use fake clock if you
-   * want to control time by manually incrementing it to test log cleanup.
+   * want to control time by manually incrementing it to test log clean.
    */
   def createTracker(
       setCheckpointDir: Boolean = true,
@@ -254,7 +254,7 @@ class ReceivedBlockTrackerSuite
     BatchAllocationEvent(time, AllocatedBlocks(Map((streamId -> blockInfos))))
   }
 
-  /** Create batch cleanup object from the given info */
+  /** Create batch clean object from the given info */
   def createBatchCleanup(time: Long, moreTimes: Long*): BatchCleanupEvent = {
     BatchCleanupEvent((Seq(time) ++ moreTimes).map(Time.apply))
   }