make CollectContextService jobSearchContextId aware

added more tests for CollectContextService
crate · Mar 2, 2015 · a22f0bd · a22f0bd
1 parent fda00a3
commit a22f0bd
Show file tree

Hide file tree

Showing 29 changed files with 982 additions and 385 deletions.
diff --git a/sql/build.gradle b/sql/build.gradle
@@ -19,6 +19,13 @@ dependencies {
     testCompile project(':testing')
     testCompile 'org.skyscreamer:jsonassert:1.2.0'
     testCompile 'com.h2database:h2:1.3.173'
+    testCompile ('org.powermock:powermock-module-junit4:1.6.1') {
+        exclude group: 'junit', module: 'junit'
+    }
+    testCompile ('org.powermock:powermock-api-mockito:1.6.1') {
+        exclude group: 'junit', module: 'junit'
+        exclude group: 'org.mockito', module: 'mockito-all'
+    }
 }
 
 buildscript {

diff --git a/sql/src/main/java/io/crate/action/sql/query/CrateSearchContext.java b/sql/src/main/java/io/crate/action/sql/query/CrateSearchContext.java
@@ -22,12 +22,15 @@
 package io.crate.action.sql.query;
 
 import com.google.common.base.Optional;
+import io.crate.Constants;
 import org.apache.lucene.util.Counter;
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.cache.recycler.CacheRecycler;
 import org.elasticsearch.cache.recycler.PageCacheRecycler;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.lease.Releasables;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.index.engine.Engine;
 import org.elasticsearch.index.service.IndexService;
@@ -43,9 +46,10 @@
 
 public class CrateSearchContext extends DefaultSearchContext {
 
+    private final Engine.Searcher engineSearcher;
+    private volatile boolean isEngineSearcherShared = false;
+
     public CrateSearchContext(long id,
-                              final int numShards,
-                              final String[] types,
                               final long nowInMillis,
                               SearchShardTarget shardTarget,
                               Engine.Searcher engineSearcher,
@@ -58,30 +62,52 @@ public CrateSearchContext(long id,
                               Counter timeEstimateCounter,
                               Optional<Scroll> scroll,
                               long keepAlive) {
-        super(id, new CrateSearchShardRequest(numShards, types, nowInMillis, scroll, indexShard),
+        super(id, new CrateSearchShardRequest(nowInMillis, scroll, indexShard),
                 shardTarget, engineSearcher, indexService,
                 indexShard, scriptService, cacheRecycler, pageCacheRecycler,
                 bigArrays, timeEstimateCounter);
+        this.engineSearcher = engineSearcher;
         if (scroll.isPresent()) {
             scroll(scroll.get());
         }
         keepAlive(keepAlive);
     }
 
+    public Engine.Searcher engineSearcher() {
+        return engineSearcher;
+    }
+
+    public void sharedEngineSearcher(boolean isShared) {
+        isEngineSearcherShared = isShared;
+    }
+
+    public boolean isEngineSearcherShared() {
+        return isEngineSearcherShared;
+    }
+
+    @Override
+    public void doClose() throws ElasticsearchException {
+        if (scanContext() != null) {
+            scanContext().clear();
+        }
+        // clear and scope phase we have
+        Releasables.close(searcher());
+        if (!isEngineSearcherShared) {
+            Releasables.close(engineSearcher);
+        }
+    }
+
+
     private static class CrateSearchShardRequest implements ShardSearchRequest {
 
-        private final int numShards;
-        private final String[] types;
+        private final String[] types = new String[]{Constants.DEFAULT_MAPPING_TYPE};
         private final long nowInMillis;
         private final Scroll scroll;
         private final String index;
         private final int shardId;
 
-        private CrateSearchShardRequest(int numShards, String[] types,
-                                        long nowInMillis, Optional<Scroll> scroll,
+        private CrateSearchShardRequest(long nowInMillis, Optional<Scroll> scroll,
                                         IndexShard indexShard) {
-            this.numShards = numShards;
-            this.types = types;
             this.nowInMillis = nowInMillis;
             this.scroll = scroll.orNull();
             this.index = indexShard.indexService().index().name();
@@ -121,7 +147,7 @@ public BytesReference extraSource() {
 
         @Override
         public int numberOfShards() {
-            return numShards;
+            return 0;
         }
 
         @Override

diff --git a/sql/src/main/java/io/crate/action/sql/query/CrateSearchService.java b/sql/src/main/java/io/crate/action/sql/query/CrateSearchService.java
@@ -24,7 +24,6 @@
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableMap;
-import io.crate.Constants;
 import io.crate.core.StringUtils;
 import io.crate.executor.transport.task.elasticsearch.SortOrder;
 import io.crate.lucene.LuceneQueryBuilder;
@@ -217,8 +216,6 @@ private SearchContext createContext(QueryShardRequest request, @Nullable Engine.
         }
         SearchContext context = new CrateSearchContext(
                 idGenerator.incrementAndGet(),
-                0, // TODO: is this necessary? seems not, wasn't set before
-                new String[] { Constants.DEFAULT_MAPPING_TYPE },
                 System.currentTimeMillis(),
                 searchShardTarget,
                 engineSearcher,

diff --git a/sql/src/main/java/io/crate/executor/transport/CollectContextService.java b/sql/src/main/java/io/crate/executor/transport/CollectContextService.java
diff --git a/sql/src/main/java/io/crate/executor/transport/TransportExecutor.java b/sql/src/main/java/io/crate/executor/transport/TransportExecutor.java
@@ -35,6 +35,7 @@
 import io.crate.metadata.Functions;
 import io.crate.metadata.ReferenceResolver;
 import io.crate.operation.ImplementationSymbolVisitor;
+import io.crate.operation.collect.CollectContextService;
 import io.crate.operation.collect.HandlerSideDataCollectOperation;
 import io.crate.operation.collect.StatsTables;
 import io.crate.operation.projectors.ProjectionToProjectorVisitor;

diff --git a/sql/src/main/java/io/crate/metadata/Routing.java b/sql/src/main/java/io/crate/metadata/Routing.java
@@ -57,7 +57,9 @@ public int numShards(String nodeId) {
             Map<String, List<Integer>> nodeRouting = locations.get(nodeId);
             if (nodeRouting != null) {
                 for (List<Integer> shardIds : nodeRouting.values()) {
-                    count += shardIds.size();
+                    if (shardIds != null) {
+                        count += shardIds.size();
+                    }
                 }
             }
         }
@@ -75,7 +77,9 @@ public int numShards() {
                 for (Map<String, List<Integer>> nodeRouting : locations.values()) {
                     if (nodeRouting != null) {
                         for (List<Integer> shardIds : nodeRouting.values()) {
-                            count += shardIds.size();
+                            if (shardIds != null) {
+                                count += shardIds.size();
+                            }
                         }
                     }
                 }