Merge pull request #107 from metamx/aggregator-close

add close method to aggregators

common/src/main/java/com/metamx/druid/aggregation/Aggregator.java

@@ -19,8 +19,6 @@
 
 package com.metamx.druid.aggregation;
 
-import java.util.Comparator;
-
 /**
  * An Aggregator is an object that can aggregate metrics.  Its aggregation-related methods (namely, aggregate() and get())
  * do not take any arguments as the assumption is that the Aggregator was given something in its constructor that
@@ -40,4 +38,5 @@ public interface Aggregator {
   Object get();
   float getFloat();
   String getName();
+  void close();
 }

common/src/main/java/com/metamx/druid/aggregation/BufferAggregator.java

@@ -88,4 +88,9 @@ public interface BufferAggregator
    * @return the float representation of the aggregate
    */
   float getFloat(ByteBuffer buf, int position);
+
+  /**
+   * Release any resources used by the aggregator
+   */
+  void close();
 }

common/src/main/java/com/metamx/druid/aggregation/CountAggregator.java

@@ -75,4 +75,10 @@ public class CountAggregator implements Aggregator
   {
     return new CountAggregator(name);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/CountBufferAggregator.java

@@ -49,4 +49,10 @@ public class CountBufferAggregator implements BufferAggregator
   {
     return buf.getLong(position);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/DoubleSumAggregator.java

@@ -89,4 +89,10 @@ public class DoubleSumAggregator implements Aggregator
   {
     return new DoubleSumAggregator(name, selector);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/DoubleSumBufferAggregator.java

@@ -59,4 +59,10 @@ public class DoubleSumBufferAggregator implements BufferAggregator
   {
     return (float) buf.getDouble(position);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/HistogramAggregator.java

@@ -80,4 +80,10 @@ public class HistogramAggregator implements Aggregator
   {
     return name;
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/HistogramBufferAggregator.java

@@ -90,4 +90,10 @@ public class HistogramBufferAggregator implements BufferAggregator
   {
     throw new UnsupportedOperationException("HistogramBufferAggregator does not support getFloat()");
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/JavaScriptAggregator.java

@@ -21,6 +21,7 @@ package com.metamx.druid.aggregation;
 
 import com.google.common.collect.Lists;
 import com.metamx.druid.processing.FloatMetricSelector;
+import org.mozilla.javascript.Context;
 
 import java.util.List;
 
@@ -33,6 +34,8 @@ public class JavaScriptAggregator implements Aggregator
     public double combine(double a, double b);
 
     public double reset();
+
+    public void close();
   }
 
   private final String name;
@@ -79,4 +82,10 @@ public class JavaScriptAggregator implements Aggregator
   {
     return name;
   }
+
+  @Override
+  public void close()
+  {
+    script.close();
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/JavaScriptAggregatorFactory.java

@@ -30,6 +30,8 @@ import com.metamx.druid.processing.MetricSelectorFactory;
 
 
 import org.mozilla.javascript.Context;
+import org.mozilla.javascript.ContextAction;
+import org.mozilla.javascript.ContextFactory;
 import org.mozilla.javascript.Function;
 import org.mozilla.javascript.Script;
 import org.mozilla.javascript.ScriptableObject;
@@ -49,7 +51,7 @@ public class JavaScriptAggregatorFactory implements AggregatorFactory
   private final List<String> fieldNames;
   private final String script;
 
-  private final JavaScriptAggregator.ScriptAggregator combiner;
+  private final JavaScriptAggregator.ScriptAggregator compiledScript;
 
   @JsonCreator
   public JavaScriptAggregatorFactory(
@@ -61,7 +63,7 @@ public class JavaScriptAggregatorFactory implements AggregatorFactory
     this.name = name;
     this.script = expression;
     this.fieldNames = fieldNames;
-    this.combiner = compileScript(script);
+    this.compiledScript = compileScript(script);
   }
 
   @Override
@@ -77,7 +79,7 @@ public class JavaScriptAggregatorFactory implements AggregatorFactory
               public FloatMetricSelector apply(@Nullable String s) { return metricFactory.makeFloatMetricSelector(s); }
             }
         ),
-        compileScript(script)
+        compiledScript
     );
   }
 
@@ -96,7 +98,7 @@ public class JavaScriptAggregatorFactory implements AggregatorFactory
               }
             }
         ),
-        compileScript(script)
+        compiledScript
     );
   }
 
@@ -109,7 +111,7 @@ public class JavaScriptAggregatorFactory implements AggregatorFactory
   @Override
   public Object combine(Object lhs, Object rhs)
   {
-    return combiner.combine(((Number) lhs).doubleValue(), ((Number) rhs).doubleValue());
+    return compiledScript.combine(((Number) lhs).doubleValue(), ((Number) rhs).doubleValue());
   }
 
   @Override
@@ -187,7 +189,7 @@ public class JavaScriptAggregatorFactory implements AggregatorFactory
   @Override
   public Object getAggregatorStartValue()
   {
-    return combiner.reset();
+    return compiledScript.reset();
   }
 
   @Override
@@ -212,22 +214,29 @@ public class JavaScriptAggregatorFactory implements AggregatorFactory
 
   public static JavaScriptAggregator.ScriptAggregator compileScript(final String script)
   {
-    final Context cx = Context.enter();
-    cx.setOptimizationLevel(9);
-    final ScriptableObject scope = cx.initStandardObjects();
+    final ContextFactory contextFactory = ContextFactory.getGlobal();
+    Context context = contextFactory.enterContext();
+    context.setOptimizationLevel(9);
 
-    Script compiledScript = cx.compileString(script, "script", 1, null);
-    compiledScript.exec(cx, scope);
+    final ScriptableObject scope = context.initStandardObjects();
+
+    Script compiledScript = context.compileString(script, "script", 1, null);
+    compiledScript.exec(context, scope);
 
     final Function fnAggregate = getScriptFunction("aggregate", scope);
     final Function fnReset = getScriptFunction("reset", scope);
     final Function fnCombine = getScriptFunction("combine", scope);
+    Context.exit();
+
 
     return new JavaScriptAggregator.ScriptAggregator()
     {
       @Override
-      public double aggregate(double current, FloatMetricSelector[] selectorList)
+      public double aggregate(final double current, final FloatMetricSelector[] selectorList)
       {
+        Context cx = Context.getCurrentContext();
+        if(cx == null) cx = contextFactory.enterContext();
+
         final int size = selectorList.length;
         final Object[] args = new Object[size + 1];
 
@@ -237,29 +246,47 @@ public class JavaScriptAggregatorFactory implements AggregatorFactory
           args[i + 1] = selectorList[i++].get();
         }
 
-        Object res = fnAggregate.call(cx, scope, scope, args);
+        final Object res = fnAggregate.call(cx, scope, scope, args);
         return Context.toNumber(res);
       }
 
       @Override
-      public double combine(double a, double b)
+      public double combine(final double a, final double b)
       {
-        Object res = fnCombine.call(cx, scope, scope, new Object[]{a, b});
+        final Object res = contextFactory.call(
+            new ContextAction()
+            {
+              @Override
+              public Object run(final Context cx)
+              {
+                return fnCombine.call(cx, scope, scope, new Object[]{a, b});
+              }
+            }
+        );
         return Context.toNumber(res);
       }
 
       @Override
       public double reset()
       {
-        Object res = fnReset.call(cx, scope, scope, new Object[]{});
+        final Object res = contextFactory.call(
+            new ContextAction()
+            {
+              @Override
+              public Object run(final Context cx)
+              {
+                return fnReset.call(cx, scope, scope, new Object[]{});
+              }
+            }
+        );
         return Context.toNumber(res);
       }
 
       @Override
-      protected void finalize() throws Throwable
-      {
-        cx.exit();
-        super.finalize();
+      public void close() {
+        if(Context.getCurrentContext() != null) {
+          Context.exit();
+        }
       }
     };
   }

common/src/main/java/com/metamx/druid/aggregation/JavaScriptBufferAggregator.java

@@ -62,4 +62,9 @@ public class JavaScriptBufferAggregator implements BufferAggregator
   {
     return (float)buf.getDouble(position);
   }
+
+  @Override
+  public void close() {
+    script.close();
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/LongSumAggregator.java

@@ -89,4 +89,10 @@ public class LongSumAggregator implements Aggregator
   {
     return new LongSumAggregator(name, selector);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/LongSumBufferAggregator.java

@@ -59,4 +59,10 @@ public class LongSumBufferAggregator implements BufferAggregator
   {
     return (float) buf.getLong(position);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/MaxAggregator.java

@@ -82,4 +82,10 @@ public class MaxAggregator implements Aggregator
   {
     return new MaxAggregator(name, selector);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/MaxBufferAggregator.java

@@ -57,4 +57,10 @@ public class MaxBufferAggregator implements BufferAggregator
   {
     return (float) buf.getDouble(position);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/MinAggregator.java

@@ -82,4 +82,10 @@ public class MinAggregator implements Aggregator
   {
     return new MinAggregator(name, selector);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/MinBufferAggregator.java

@@ -57,4 +57,10 @@ public class MinBufferAggregator implements BufferAggregator
   {
     return (float) buf.getDouble(position);
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/NoopAggregator.java

@@ -59,4 +59,10 @@ public class NoopAggregator implements Aggregator
   {
     return name;
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/aggregation/NoopBufferAggregator.java

@@ -46,4 +46,10 @@ public class NoopBufferAggregator implements BufferAggregator
   {
     return 0;
   }
+
+  @Override
+  public void close()
+  {
+    // no resources to cleanup
+  }
 }

common/src/main/java/com/metamx/druid/processing/MetricSelectorUtils.java

@@ -0,0 +1,5 @@
+package com.metamx.druid.processing;
+
+public class MetricSelectorUtils
+{
+}

server/src/main/java/com/metamx/druid/query/group/GroupByQueryEngine.java

@@ -31,6 +31,8 @@ import com.metamx.common.ISE;
 import com.metamx.common.guava.BaseSequence;
 import com.metamx.common.guava.FunctionalIterator;
 import com.metamx.common.guava.Sequence;
+import com.metamx.common.guava.Sequences;
+import com.metamx.common.parsers.CloseableIterator;
 import com.metamx.druid.StorageAdapter;
 import com.metamx.druid.aggregation.AggregatorFactory;
 import com.metamx.druid.aggregation.BufferAggregator;
@@ -47,6 +49,7 @@ import org.joda.time.DateTime;
 import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
+import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Iterator;
 import java.util.List;
@@ -85,33 +88,47 @@ public class GroupByQueryEngine
 
     final ResourceHolder<ByteBuffer> bufferHolder = intermediateResultsBufferPool.take();
 
-    return new BaseSequence<Row, Iterator<Row>>(
-        new BaseSequence.IteratorMaker<Row, Iterator<Row>>()
+    return Sequences.concat(
+        new BaseSequence<Sequence<Row>, Iterator<Sequence<Row>>>(new BaseSequence.IteratorMaker<Sequence<Row>, Iterator<Sequence<Row>>>()
         {
           @Override
-          public Iterator<Row> make()
+          public Iterator<Sequence<Row>> make()
           {
             return FunctionalIterator
-                .create(cursors.iterator())
-                .transformCat(
-                    new Function<Cursor, Iterator<Row>>()
-                    {
-                      @Override
-                      public Iterator<Row> apply(@Nullable final Cursor cursor)
-                      {
-                        return new RowIterator(query, cursor, bufferHolder.get());
-                      }
-                    }
-                );
+              .create(cursors.iterator())
+              .transform(new Function<Cursor, Sequence<Row>>()
+                        {
+                          @Override
+                          public Sequence<Row> apply(@Nullable final Cursor cursor)
+                          {
+                            return new BaseSequence<Row, CloseableIterator<Row>>(
+                                new BaseSequence.IteratorMaker<Row, CloseableIterator<Row>>()
+                                {
+                                  @Override
+                                  public  CloseableIterator<Row> make()
+                                  {
+                                    return new RowIterator(query, cursor, bufferHolder.get());
+                                  }
+
+                                  @Override
+                                  public void cleanup(CloseableIterator iterFromMake)
+                                  {
+                                    Closeables.closeQuietly(iterFromMake);
+                                  }
+                                }
+                            );
+                          }
+              });
           }
 
           @Override
-          public void cleanup(Iterator<Row> iterFromMake)
+          public void cleanup(Iterator<Sequence<Row>> iterFromMake)
           {
             Closeables.closeQuietly(bufferHolder);
           }
-        }
+        })
     );
+
   }
 
   private static class RowUpdater
@@ -248,7 +265,7 @@ public class GroupByQueryEngine
     }
   }
 
-  private class RowIterator implements Iterator<Row>
+  private class RowIterator implements CloseableIterator<Row>
   {
     private final GroupByQuery query;
     private final Cursor cursor;
@@ -384,5 +401,12 @@ public class GroupByQueryEngine
     {
       throw new UnsupportedOperationException();
     }
+
+    public void close() {
+      // cleanup
+      for(BufferAggregator agg : aggregators) {
+        agg.close();
+      }
+    }
   }
 }

server/src/main/java/com/metamx/druid/query/timeseries/TimeseriesQueryEngine.java

@@ -80,7 +80,14 @@ public class TimeseriesQueryEngine
                       bob.addMetric(postAgg);
                     }
 
-                    return bob.build();
+                    Result<TimeseriesResultValue> retVal = bob.build();
+
+                    // cleanup
+                    for(Aggregator agg : aggregators) {
+                      agg.close();
+                    }
+
+                    return retVal;
                   }
                 }
             ).iterator();