ARROW-6199: [Java] Avro adapter avoid potential resource leak.

tianchen92 · emkornfield · commit dd4532a0cdac · 2019-08-15T19:49:53.000-07:00
Related to [ARROW-6199](https://issues.apache.org/jira/browse/ARROW-6199). Currently, avro consumer interface has no close API, which may cause resource leak like AvroBytesConsumer#cacheBuffer. To resolve this, make consumer extends AutoCloseable and create CompositeAvroConsumer to encompasses consume and close logic. Closes #5059 from tianchen92/ARROW-6199 and squashes the following commits: d60d94c <tianchen> fix 42f22da <tianchen> clear vectors in close 5b91da7 <tianchen> fix comments 3ffc076 <tianchen> ARROW-6199: Avro adapter avoid potential resource leak. Authored-by: tianchen <niki.lj@alibaba-inc.com> Signed-off-by: Micah Kornfield <emkornfield@gmail.com>
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/AvroToArrowUtils.java b/java/adapter/avro/src/main/java/org/apache/arrow/AvroToArrowUtils.java
@@ -20,7 +20,6 @@
 import static org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE;
 import static org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE;
 
-import java.io.EOFException;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -37,6 +36,7 @@
 import org.apache.arrow.consumers.AvroNullConsumer;
 import org.apache.arrow.consumers.AvroStringConsumer;
 import org.apache.arrow.consumers.AvroUnionsConsumer;
+import org.apache.arrow.consumers.CompositeAvroConsumer;
 import org.apache.arrow.consumers.Consumer;
 import org.apache.arrow.consumers.NullableTypeConsumer;
 import org.apache.arrow.memory.BufferAllocator;
@@ -246,19 +246,15 @@ public static VectorSchemaRoot avroToArrowVectors(Schema schema, Decoder decoder
 
     VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, 0);
 
-    int valueCount = 0;
-    while (true) {
-      try {
-        for (Consumer consumer : consumers) {
-          consumer.consume(decoder);
-        }
-        valueCount++;
-        //reach end will throw EOFException.
-      } catch (EOFException eofException) {
-        root.setRowCount(valueCount);
-        break;
-      }
+    CompositeAvroConsumer compositeConsumer = null;
+    try {
+      compositeConsumer = new CompositeAvroConsumer(consumers);
+      compositeConsumer.consume(decoder, root);
+    } catch (Exception e) {
+      compositeConsumer.close();
+      throw new RuntimeException("Error occurs while consume process.", e);
     }
+
     return root;
   }
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroBooleanConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroBooleanConsumer.java
@@ -63,4 +63,9 @@ public FieldVector getVector() {
     return this.vector;
   }
 
+  @Override
+  public void close() throws Exception {
+    writer.close();
+  }
+
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroBytesConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroBytesConsumer.java
@@ -79,4 +79,9 @@ public void setPosition(int index) {
   public FieldVector getVector() {
     return vector;
   }
+
+  @Override
+  public void close() throws Exception {
+    writer.close();
+  }
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroDoubleConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroDoubleConsumer.java
@@ -62,4 +62,9 @@ public void setPosition(int index) {
   public FieldVector getVector() {
     return vector;
   }
+
+  @Override
+  public void close() throws Exception {
+    writer.close();
+  }
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroFloatConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroFloatConsumer.java
@@ -62,4 +62,9 @@ public void setPosition(int index) {
   public FieldVector getVector() {
     return this.vector;
   }
+
+  @Override
+  public void close() throws Exception {
+    writer.close();
+  }
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroIntConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroIntConsumer.java
@@ -62,4 +62,9 @@ public void setPosition(int index) {
   public FieldVector getVector() {
     return this.vector;
   }
+
+  @Override
+  public void close() throws Exception {
+    writer.close();
+  }
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroLongConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroLongConsumer.java
@@ -62,4 +62,9 @@ public void setPosition(int index) {
   public FieldVector getVector() {
     return this.vector;
   }
+
+  @Override
+  public void close() throws Exception {
+    writer.close();
+  }
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroNullConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroNullConsumer.java
@@ -48,4 +48,9 @@ public void setPosition(int index) {}
   public FieldVector getVector() {
     return this.vector;
   }
+
+  @Override
+  public void close() {
+    vector.close();
+  }
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroStringConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroStringConsumer.java
@@ -80,4 +80,9 @@ public void setPosition(int index) {
   public FieldVector getVector() {
     return this.vector;
   }
+
+  @Override
+  public void close() throws Exception {
+    writer.close();
+  }
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroUnionsConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/AvroUnionsConsumer.java
@@ -31,7 +31,7 @@
  */
 public class AvroUnionsConsumer implements Consumer {
 
-  private Consumer[] indexDelegates;
+  private Consumer[] delegates;
   private Types.MinorType[] types;
 
   private UnionWriter writer;
@@ -40,11 +40,11 @@ public class AvroUnionsConsumer implements Consumer {
   /**
    * Instantiate a AvroUnionConsumer.
    */
-  public AvroUnionsConsumer(UnionVector vector, Consumer[] indexDelegates, Types.MinorType[] types) {
+  public AvroUnionsConsumer(UnionVector vector, Consumer[] delegates, Types.MinorType[] types) {
 
     this.writer = new UnionWriter(vector);
     this.vector = vector;
-    this.indexDelegates = indexDelegates;
+    this.delegates = delegates;
     this.types = types;
   }
 
@@ -53,7 +53,7 @@ public void consume(Decoder decoder) throws IOException {
     int fieldIndex = decoder.readInt();
     int position = writer.getPosition();
 
-    Consumer delegate = indexDelegates[fieldIndex];
+    Consumer delegate = delegates[fieldIndex];
 
     vector.setType(position, types[fieldIndex]);
     // In UnionVector we need to set sub vector writer position before consume a value
@@ -80,4 +80,12 @@ public FieldVector getVector() {
     vector.setValueCount(writer.getPosition());
     return this.vector;
   }
+
+  @Override
+  public void close() throws Exception {
+    writer.close();
+    for (Consumer delegate: delegates) {
+      delegate.close();
+    }
+  }
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/CompositeAvroConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/CompositeAvroConsumer.java
@@ -0,0 +1,69 @@
+/*
+ * 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.arrow.consumers;
+
+import java.io.EOFException;
+import java.io.IOException;
+import java.util.List;
+
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.avro.io.Decoder;
+
+/**
+ * Composite consumer which hold all consumers.
+ * It manages the consume and cleanup process.
+ */
+public class CompositeAvroConsumer implements AutoCloseable {
+
+  private final List<Consumer> consumers;
+
+  public CompositeAvroConsumer(List<Consumer> consumers) {
+    this.consumers = consumers;
+  }
+
+  /**
+   * Consume decoder data and write into {@link VectorSchemaRoot}.
+   */
+  public void consume(Decoder decoder, VectorSchemaRoot root) throws IOException {
+    int valueCount = 0;
+    while (true) {
+      try {
+        for (Consumer consumer : consumers) {
+          consumer.consume(decoder);
+        }
+        valueCount++;
+        //reach end will throw EOFException.
+      } catch (EOFException eofException) {
+        root.setRowCount(valueCount);
+        break;
+      }
+    }
+  }
+
+  @Override
+  public void close() {
+    // clean up
+    for (Consumer consumer : consumers) {
+      try {
+        consumer.close();
+      } catch (Exception e) {
+        throw new RuntimeException("Error occurs in close.", e);
+      }
+    }
+  }
+}
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/Consumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/Consumer.java
@@ -25,7 +25,7 @@
 /**
  * Interface that is used to consume values from avro decoder.
  */
-public interface Consumer {
+public interface Consumer extends AutoCloseable {
 
   /**
    * Consume a specific type value from avro decoder and write it to vector.
@@ -48,4 +48,9 @@ public interface Consumer {
    * Get the vector within the consumer.
    */
   FieldVector getVector();
+
+  /**
+   * Close this consumer when occurs exception to avoid potential leak.
+   */
+  void close() throws Exception;
 }
diff --git a/java/adapter/avro/src/main/java/org/apache/arrow/consumers/NullableTypeConsumer.java b/java/adapter/avro/src/main/java/org/apache/arrow/consumers/NullableTypeConsumer.java
@@ -65,4 +65,9 @@ public FieldVector getVector() {
     return delegate.getVector();
   }
 
+  @Override
+  public void close() throws Exception {
+    delegate.close();
+  }
+
 }

Original file line number	Diff line number	Diff line change
`@@ -63,4 +63,9 @@ public FieldVector getVector() {`
`63`	`63`	`return this.vector;`
`64`	`64`	`}`
`65`	`65`
	`66`	`+ @Override`
	`67`	`+ public void close() throws Exception {`
	`68`	`+ writer.close();`
	`69`	`+ }`
	`70`	`+`
`66`	`71`	`}`
Original file line number	Diff line number	Diff line change
`@@ -79,4 +79,9 @@ public void setPosition(int index) {`
`79`	`79`	`public FieldVector getVector() {`
`80`	`80`	`return vector;`
`81`	`81`	`}`
	`82`	`+`
	`83`	`+ @Override`
	`84`	`+ public void close() throws Exception {`
	`85`	`+ writer.close();`
	`86`	`+ }`
`82`	`87`	`}`
Original file line number	Diff line number	Diff line change
`@@ -62,4 +62,9 @@ public void setPosition(int index) {`
`62`	`62`	`public FieldVector getVector() {`
`63`	`63`	`return vector;`
`64`	`64`	`}`
	`65`	`+`
	`66`	`+ @Override`
	`67`	`+ public void close() throws Exception {`
	`68`	`+ writer.close();`
	`69`	`+ }`
`65`	`70`	`}`
Original file line number	Diff line number	Diff line change
`@@ -48,4 +48,9 @@ public void setPosition(int index) {}`
`48`	`48`	`public FieldVector getVector() {`
`49`	`49`	`return this.vector;`
`50`	`50`	`}`
	`51`	`+`
	`52`	`+ @Override`
	`53`	`+ public void close() {`
	`54`	`+ vector.close();`
	`55`	`+ }`
`51`	`56`	`}`
Original file line number	Diff line number	Diff line change
`@@ -80,4 +80,9 @@ public void setPosition(int index) {`
`80`	`80`	`public FieldVector getVector() {`
`81`	`81`	`return this.vector;`
`82`	`82`	`}`
	`83`	`+`
	`84`	`+ @Override`
	`85`	`+ public void close() throws Exception {`
	`86`	`+ writer.close();`
	`87`	`+ }`
`83`	`88`	`}`