ElasticSearch：剖析query_and_fetch和query_then_fetch的區別

在ElasticSearch實際使用中遇到了一個問題，就是在數據量很大的情況下做聚合查詢（aggregation）會導致內存溢出。當時看了文檔，猜測修改search_type能避免內存溢出。實際測試發現，在數據量相同的情況下，search_type設置爲query_and_fetch的聚合查詢不會導致內存溢出，而默認的query_then_fetch則會內存溢出。本文就從源碼層面分析這兩種search_type的區別。

搜索請求處理過程

當一個搜索請求發送到節點之後，節點首先判斷出這是一個搜索請求，然後將這個請求傳遞給TransportSearchAction。這個類負責處理所有的搜索請求。

TransportSearchAction負責根據搜索請求中的search_type將本次的搜索請求傳遞給對應的類。核心代碼如下：

if (searchRequest.searchType() == DFS_QUERY_THEN_FETCH) {
    dfsQueryThenFetchAction.execute(searchRequest, listener);
} else if (searchRequest.searchType() == SearchType.QUERY_THEN_FETCH) {
    queryThenFetchAction.execute(searchRequest, listener);
} else if (searchRequest.searchType() == SearchType.DFS_QUERY_AND_FETCH) {
    dfsQueryAndFetchAction.execute(searchRequest, listener);
} else if (searchRequest.searchType() == SearchType.QUERY_AND_FETCH) {
    queryAndFetchAction.execute(searchRequest, listener);
} else if (searchRequest.searchType() == SearchType.SCAN) {
    scanAction.execute(searchRequest, listener);
} else if (searchRequest.searchType() == SearchType.COUNT) {
    countAction.execute(searchRequest, listener);
}

query_and_fetch執行過程

query_and_fetch搜索對應的類爲TransportSearchQueryAndFetchAction。它在執行的時候會啓動一個異步任務，對應的代碼如下：

@Override
protected void doExecute(SearchRequest searchRequest, ActionListener<SearchResponse> listener) {
    new AsyncAction(searchRequest, listener).start();
}

AsyncAction的入口方法是start，定義在它的基類BaseAsyncAction中。

BaseAsyncAction中的start方法首先確定有哪些分片需要處理，然後給每個需要處理的分片都啓動一個異步的搜索任務，對應的關鍵代碼如下：

for (final ShardIterator shardIt : shardsIts) {
    ...
    performFirstPhase(shardIndex, shardIt, shard);
    ...
}

performFirstPhase的作用就是生成一個內部的分片搜索請求，這種請求只針對一個分片，然後調用了一個子類的方法sendExecuteFirstPhase，讓子類選擇的處理方式。對應的代碼如下：

sendExecuteFirstPhase(node, firstRequest, new SearchServiceListener<FirstResult>() {
    @Override
    public void onResult(FirstResult result) {
        onFirstPhaseResult(shardIndex, shard, result, shardIt);
    }
    ...
});

onFirstPhaseResult主要作用是調用子類的moveToSecondPhase。這個方法在executeFetchPhase之後才執行的，因此在其後面再介紹。

sendExecuteFirstPhase方法是抽象的。在搜索模式爲query_and_fetch時，對分片請求的處理方式是調用sendExecuteFetch。子類對它的實現代碼如下：

@Override
protected void sendExecuteFirstPhase(DiscoveryNode node, ShardSearchRequest request, SearchServiceListener<QueryFetchSearchResult> listener) {
    searchService.sendExecuteFetch(node, request, listener);
}

sendExecuteFetch定義在SearchServiceTransportAction中，它會將分片搜索請求轉發到對應的節點上。當然，如果對應的節點是自己這個節點，就不轉發了，直接執行executeFetchPhase。

executeFetchPhase的執行過程

executeFetchPhase中，首先初始化一個SearchContext，然後調用queryPhase，然後調用fetchPhase，最後清理SearchContext。對應的關鍵代碼如下：

public QueryFetchSearchResult executeFetchPhase(ShardSearchRequest request) throws ElasticsearchException {
    // 初始化SearchContext
    final SearchContext context = createAndPutContext(request);
    ...

    try {
        // 執行queryPhase
        ...
        queryPhase.execute(context);
        ...

        // 執行fetchPhase
        ...
        fetchPhase.execute(context);
        ...

        if (context.scroll() == null) {
           freeContext(context.id());
       }

       ...

        // 返回結果
        return new QueryFetchSearchResult(context.queryResult(), context.fetchResult());
    } catch (Throwable e) {
        ...
    } finally {
        // 清理SearchContext
        cleanContext(context);
    }
}

從上面的代碼可以看出，queryPhase和fetchPhase是連續執行的，這就是query_and_fetch的含義。當這兩個階段執行完成之後，如果不是滾動查詢，就直接調用freeContext將SearchContext從activeContexts中刪除。

freeContext和clearContext有什麼區別呢？

freeContext的主要作用是從activeContexts中刪掉指定的SearchContext，這樣JVM就能將這部分的內存進行回收了。相關的代碼如下：

public boolean freeContext(long id) {
    final SearchContext context = activeContexts.remove(id);
    ...
    context.close();
    ...
}

cleanContext的主要作用是釋放部分內存空間，然後將指定的context從ThreadLocal中刪掉。相關的代碼如下：

private void cleanContext(SearchContext context) {
    context.clearReleasables(Lifetime.PHASE);
    SearchContext.removeCurrent();
}

SearchContext.clearReleasables的主要作用是將可以釋放的內存釋放掉。那怎樣判斷哪些資源是可以釋放的呢？這個需要別的類在裏面進行註冊了。註冊的時候調用addReleasable，告知哪個階段可以清理哪些資源。相關的代碼如下：

private Multimap<Lifetime, Releasable> clearables = null;
public void addReleasable(Releasable releasable, Lifetime lifetime) {
    ...
    clearables.put(lifetime, releasable);
}

public void clearReleasables(Lifetime lifetime) {
    ...
    List<Collection<Releasable>> releasables = new ArrayList<>();
    for (Lifetime lc : Lifetime.values()) {
        if (lc.compareTo(lifetime) > 0) {
            break;
        }
        releasables.add(clearables.removeAll(lc));
    }
    Releasables.close(Iterables.concat(releasables));
    ...
}

上面這段代碼中的for循環是爲了收集指定階段，以及該階段之前所有可清理的資源。因此clearReleasables的作用是清理掉指定階段以及該階段之前的所有可清理的資源。

SearchContext.removeCurrent代碼如下。主要就是將當前的SearchContext刪除掉。

private static ThreadLocal<SearchContext> current = new ThreadLocal<>();
public static void removeCurrent() {
    current.remove();
    QueryParseContext.removeTypes();
}

分析一下executeFetchPhase返回結果。從代碼中可以看出，返回的結果中包含query的結果和fetch的結果。對應的兩個類分別是QuerySearchResult和FetchSearchResult。

QuerySearchResult

QuerySearchResult定義如下：

public class QuerySearchResult extends TransportResponse implements QuerySearchResultProvider {

    private long id;
    private SearchShardTarget shardTarget;
    private int from;
    private int size;
    private TopDocs topDocs;
    private InternalFacets facets;
    private InternalAggregations aggregations;
    private Suggest suggest;
    private boolean searchTimedOut;

    ...
}

其中最重要的應該是topDocs和aggregations。

再看看TopDocs的定義：

public class TopDocs {
  public int totalHits;
  public ScoreDoc[] scoreDocs;
  private float maxScore;
  ...
}  

裏面包含了搜索命中的結果數量、文檔編號、文檔匹配分數、最大匹配分數。再看看ScoreDoc是如何定義的。

public class ScoreDoc {
  public float score;
  public int doc;
  public int shardIndex;
  ...
}

這裏的doc就是內部的文檔編號，可以通過IndexSearcher#doc(int)方法獲取對應的文檔內容。

因此QuerySearchResult中只包含了內部的文檔編號、文檔的匹配分值。

FetchSearchResult

FetchSearchResult的定義如下：

public class FetchSearchResult extends TransportResponse implements FetchSearchResultProvider {
    private long id;
    private SearchShardTarget shardTarget;
    private InternalSearchHits hits;
    private transient int counter;
    ...
}

它包含了InternalSearchHits。InternalSearchHits的定義如下：

public class InternalSearchHits implements SearchHits {
    private InternalSearchHit[] hits;
    public long totalHits;
    private float maxScore;
    ...
}

它包含了InternalSearchHit。InternalSearchHit的定義如下：

public class InternalSearchHit implements SearchHit {
    ...
    private Map<String, Object> sourceAsMap;
    private byte[] sourceAsBytes;
    ...
}

它包含了文檔的原始內容和解析後的內容。

因此FetchSearchResult包含了文檔的具體內容。

moveToSecondPhase執行過程

在上面的executeFetchPhase執行完成之後，得到query結果和fetch結果之後，就執行moveToSecondPhase了。

moveToSecondPhase在BaseAsyncAction中是一個抽象方法，它在子類TransportSearchQueryAndFetchAction中的定義如下：

@Override
protected void moveToSecondPhase() throws Exception {
    threadPool.executor(ThreadPool.Names.SEARCH).execute(new Runnable() {
        @Override
        public void run() {
            ...
            sortedShardList = searchPhaseController.sortDocs(useScroll, firstResults);
            final InternalSearchResponse internalResponse = searchPhaseController.merge(sortedShardList, firstResults, firstResults);
            ...
            listener.onResponse(new SearchResponse(internalResponse, scrollId, expectedSuccessfulOps, successfulOps.get(), buildTookInMillis(), buildShardFailures()));
        }
    });
}

從代碼中可以看出，搜索的第二階段是在search線程池中提交一個任務，首先是對分片結果進行整體排序，然後將搜索結果進行合併。這裏面分別調用了searchPhaseController.sortDocs和searchPhaseController.merge兩個方法。

首先看看sortDocs做了什麼。它的關鍵代碼如下：

public ScoreDoc[] sortDocs(boolean scrollSort, AtomicArray<? extends QuerySearchResultProvider> resultsArr) throws IOException {
    ...
    TopDocs mergedTopDocs = TopDocs.merge(sort, from, topN, shardTopDocs);
    return mergedTopDocs.scoreDocs;
}

從代碼中可以看出，對文檔進行排序的時候調用了TopDocs.merge方法。這個方法的關鍵代碼如下：

public static TopDocs merge(Sort sort, int start, int size, TopDocs[] shardHits) throws IOException {
    final PriorityQueue<ShardRef> queue;
    ...
    queue = new ScoreMergeSortQueue(shardHits);
    ...

    for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
        ...
        queue.add(new ShardRef(shardIDX));
        ...
    }

    ...

    // hits儲存最終的結果
    hits = new ScoreDoc[Math.min(size, availHitCount - start)];

    // 每次取出最小值放到結果中
    while (hitUpto < numIterOnHits) {
        ...

        // 利用優先級隊列，從所有分片中取出分值最小的文檔，加入到hits結果中
        ShardRef ref = queue.pop();
        final ScoreDoc hit = shardHits[ref.shardIndex].scoreDocs[ref.hitIndex++];
        hit.shardIndex = ref.shardIndex;
        ...
        hits[hitUpto - start] = hit;
        ...

        hitUpto++;

        // 數組中的數據讀取完了，就從優先級隊列中排除
        if (ref.hitIndex < shardHits[ref.shardIndex].scoreDocs.length) {
            queue.add(ref);
        }
    }
}

這段代碼剛開始看起來有點複雜，看不太懂。其實它是歸併排序的變種。因爲普通的歸併排序只針對兩個數組。排序的時候每次從輸入的兩個數組中取出最小的元素，放到結果中。而這裏輸入的數組會有多個，每次也要取出最小的元素，放到結果中。如何快速的從多個數組中取出最小的元素呢？這裏利用了優先級隊列。

再看searchPhaseController.merge做了什麼呢？主要是合併facet結果、aggregation結果、hits結果、suggest結果、count結果。這裏就看一下最關鍵的hits是怎樣合併的。相關的代碼如下：

List<InternalSearchHit> hits = new ArrayList<>();
for (ScoreDoc shardDoc : sortedDocs) {
    FetchSearchResult fetchResult = ...;
    ...
    int index = fetchResult.counterGetAndIncrement();
    if (index < fetchResult.hits().internalHits().length) {
        ...
        InternalSearchHit searchHit = fetchResult.hits().internalHits()[index];
        ...
        hits.add(searchHit);
    }
    ...
}

hit的合併過程就是將fetchResult中所有文檔的具體內容組合成一個新的列表。其他的facet、aggregations、suggest、count的合併過程也是類似。

query_and_fetch下的from和size

舉個例子，如果在搜索的時候指定search_type爲query_and_fetch，再指定size爲10，那麼就會返回50個結果。這是爲什麼呢？原來，在fetch的時候就已經把from和size參數用掉了，導致每個分片都返回了10個文檔，如果有5個分片的話，那麼最後合併的結果就是50個。

如果query_and_fetch情況下這兩和參數的作用和query_then_fetch的行爲一樣。那麼一次就要取出from+size個文檔。如果from比較大，那麼取出的文檔足以撐爆內存。而如果在fetch階段就直接使用這兩個參數，每個分片最多就取出size個文檔，from稍微大一些也沒有關係。因此，這種行爲的設計也是可以理解的。

query_then_fetch執行過程

query_then_fetch的搜索邏輯在TransportSearchQueryThenFetchAction中。它在執行的時候會啓動一個異步任務。相關代碼如下：

@Override
protected void doExecute(SearchRequest searchRequest, ActionListener<SearchResponse> listener) {
    new AsyncAction(searchRequest, listener).start();
}

start方法定義在它的基類BaseAsyncTask中，主要作用是給每個需要搜索的分片單獨啓動一個異步任務，並讓子類選擇處理方式。在第一步處理完成之後，會調用子類中的moveToSecondPhase繼續執行第二階段的計算任務。在query_then_fetch中，子類的處理方式是sendExecuteQuery。相關的代碼如下：

@Override
protected void sendExecuteFirstPhase(DiscoveryNode node, ShardSearchRequest request, SearchServiceListener<QuerySearchResult> listener) {
    searchService.sendExecuteQuery(node, request, listener);
}

sendExecuteQuery的作用是什麼呢？它會將query請求發送到對應的節點上。如果對應的節點就是自己，那麼就直接調用executeQueryPhase，執行query任務。

executeQueryPhase的主要作用是初始化SearchContext，然後執行queryPhase，最後清理SearchContext。返回的結果是QuerySearchResult，它只包含文檔編號和必要的排序分值。相關的代碼如下：

public QuerySearchResult executeQueryPhase(ShardSearchRequest request) throws ElasticsearchException {
    final SearchContext context = createAndPutContext(request);
    try {
       ...
       queryPhase.execute(context);
       ...
       if (context.searchType() == SearchType.COUNT) {
           freeContext(context.id());
       }
       ...
       return context.queryResult();
    } catch(...) {
        ...
    } finally {
       ...
       cleanContext(context);
    }
}

代碼裏面可以看到一個細節，就是只有在search_type爲count的時候纔會調用freeContext。也就是說，如果search_type不是count，那麼SearchContext仍然會駐留在內存中。cleanContext和freeContext的區別在前面的文章裏面講過了。

當queryPhase結束之後，就開始第二階段了。第二階段從moveToSecondPhase開始。它的代碼定義在TransportSearchQueryThenFetchAction中。主要的作用是將第一階段獲取的文檔編號進行排序。排序完成之後再根據文檔編號獲取文檔裏面實際的內容。相關的代碼如下：

@Override
protected void moveToSecondPhase() throws Exception {
    ...
    sortedShardList = searchPhaseController.sortDocs(useScroll, firstResults);
    searchPhaseController.fillDocIdsToLoad(docIdsToLoad, sortedShardList);
    ...
    final AtomicInteger counter = new AtomicInteger(docIdsToLoad.asList().size());
    for (AtomicArray.Entry<IntArrayList> entry : docIdsToLoad.asList()) {
        QuerySearchResult queryResult = firstResults.get(entry.index);
        ...
        executeFetch(entry.index, queryResult.shardTarget(), counter, fetchSearchRequest, node);
    }
}

sortDocs前面已經講過了，作用是將文檔編號根據每個文檔的匹配分值進行排序。

executeFetch相關的代碼如下。它的作用是調用searchService開啓一個異步任務，根據文檔編號獲取文檔的具體內容，並將結果存放到fetchResults中。根據counter判斷如果所有的fetch任務都執行完了，就調用finishHim來完成本次查詢結果。

void executeFetch(final int shardIndex, final SearchShardTarget shardTarget, final AtomicInteger counter, final FetchSearchRequest fetchSearchRequest, DiscoveryNode node) {
    searchService.sendExecuteFetch(node, fetchSearchRequest, new SearchServiceListener<FetchSearchResult>() {
        @Override
        public void onResult(FetchSearchResult result) {
            ...
            fetchResults.set(shardIndex, result);
            if (counter.decrementAndGet() == 0) {
                finishHim();
            }
        }
        ...
    });
}

sendExecuteFetch最後會調用executeFetchPhase。它的關鍵代碼如下：

public FetchSearchResult executeFetchPhase(FetchSearchRequest request) throws ElasticsearchException {
    final SearchContext context = findContext(request.id());
    ...
    fetchPhase.execute(context);
    ...
    freeContext(request.id());
    ...
    return context.fetchResult();
}

從代碼中可以看出，在fetch階段結束之後纔會釋放SearchContext。

finishHim相關的代碼如下。它的作用是合併每個分片的查詢結果，讓後將合併結果通知給listener。讓它完成最後的查詢結果。searchPhaseController.merge在前面講過了，主要作用是

private void finishHim() {
    ...
    threadPool.executor(ThreadPool.Names.SEARCH).execute(new Runnable() {
        @Override
        public void run() {
            ...
            final InternalSearchResponse internalResponse = searchPhaseController.merge(sortedShardList, firstResults, fetchResults);
            ...
            listener.onResponse(new SearchResponse(internalResponse, scrollId, expectedSuccessfulOps, successfulOps.get(), buildTookInMillis(), buildShardFailures()));
            ...
        }
    });
    ...
}

搜索模式對聚合查詢的影響

query_and_fetch和query_then_fetch下的聚合查詢有什麼區別呢？首先，前面已經講過了，query_and_fetch在executeFetchPhase的時候就把SearchContext釋放掉了。

public QueryFetchSearchResult executeFetchPhase(ShardSearchRequest request) throws ElasticsearchException {
    final SearchContext context = createAndPutContext(request);
    ...
    queryPhase.execute(context);
    ...
    fetchPhase.execute(context);
    ...
    freeContext(context.id());
    ...
    return new QueryFetchSearchResult(context.queryResult(), context.fetchResult());
}

而在query_then_fetch中，query階段不會釋放SearchContext，在fetch階段結束之後纔會釋放SearchContext。兩段相關的代碼如下：

public QuerySearchResult executeQueryPhase(ShardSearchRequest request) throws ElasticsearchException {
    final SearchContext context = createAndPutContext(request);
    ...
    queryPhase.execute(context);
    ...
    return context.queryResult();
}

public FetchSearchResult executeFetchPhase(FetchSearchRequest request) throws ElasticsearchException {
    final SearchContext context = findContext(request.id());
    ...
    fetchPhase.execute(context);
    ...
    freeContext(request.id());
    ...
    return context.fetchResult();
}

在聚類搜索的時候，內存佔用比較大的是FieldData數據，這些數據儲存在SearchContext中。在query_then_fetch的搜索模式下，載入FieldData是在query階段完成的。而query階段不會釋放內存，因此內存中會存放所有分片的FieldData，從而導致內存溢出。

但是我還是有個疑問，爲什麼不在query階段結束之後立即釋放SearchContext中的FieldData呢？這樣也就不會有內存溢出的問題了。