Add max_batch_size option for bfs_executor

Author: Totktonada

graphql/bfs_executor.lua

@@ -176,6 +176,7 @@
 ---         }
 ---     }
 
+local json = require('json')
 local utils = require('graphql.utils')
 local core_util = require('graphql.core.util')
 local core_types = require('graphql.core.types')
@@ -725,14 +726,20 @@ local function invoke_resolve_list(prepared_object_list, context, opts)
     local qcontext = opts.qcontext
     local accessor = opts.accessor
     local is_item_cache_only = opts.is_item_cache_only or false
-
-    fetch_resolve_list(prepared_object_list, {accessor = accessor,
-        qcontext = qcontext})
+    local max_batch_size = opts.max_batch_size
 
     local open_set = {}
     local cache_only_open_set = {}
 
-    for _, prepared_object in ipairs(prepared_object_list) do
+    local last_fetched_object_num = 0
+    for i, prepared_object in ipairs(prepared_object_list) do
+        if i > last_fetched_object_num then
+            local _, size = fetch_resolve_list(prepared_object_list,
+                {accessor = accessor, qcontext = qcontext,
+                max_batch_size = max_batch_size, start_from = i})
+            last_fetched_object_num = last_fetched_object_num + size
+        end
+
         local child = invoke_resolve(prepared_object, context,
             {qcontext = qcontext, is_item_cache_only = is_item_cache_only})
         local child_open_set = child.open_set
@@ -753,23 +760,26 @@ end
 -- Analyze prepared requests and prefetch in batches {{{
 
 fetch_first_same = function(open_set, opts)
+    local func_name = 'bfs_executor.fetch_first_same'
     local opts = opts or {}
     local accessor = opts.accessor
     local qcontext = opts.qcontext
+    local max_batch_size = opts.max_batch_size
 
     if not accessor:cache_is_supported() then return nil, 0 end
 
     local size = 0
 
     local batches = {}
     for i, item in ipairs(open_set) do
+        if i > max_batch_size then break end
         if item.prepared_object == nil then break end
         local prepared_object = item.prepared_object
 
         for field_name, field_info in pairs(prepared_object.fields_info) do
             local prepared_resolve = field_info.prepared_resolve
             if prepared_resolve.is_calculated then
-                -- don't update size, because we don't add request to 'batches'
+                size = i
                 break
             end
             local prepared_select = prepared_resolve.prepared_select
@@ -780,7 +790,9 @@ fetch_first_same = function(open_set, opts)
             local iterator_opts = request_opts.iterator_opts
 
             if i == 1 then
-                assert(batches[field_name] == nil, 'XXX') -- XXX
+                assert(batches[field_name] == nil,
+                    ('internal error: %s: field names "%s" clash'):format(
+                    func_name, field_name))
                 batches[field_name] = {
                     collection_name = collection_name,
                     index_name = index_name,
@@ -792,10 +804,8 @@ fetch_first_same = function(open_set, opts)
                 local ok =
                     batches[field_name].collection_name == collection_name and
                     batches[field_name].index_name == index_name and
-                    utils.is_subtable(batches[field_name].iterator_opts or {},
-                        iterator_opts or {}) and
-                    utils.is_subtable(iterator_opts or {},
-                        batches[field_name].iterator_opts or {})
+                    utils.are_tables_same(batches[field_name].iterator_opts or
+                        {}, iterator_opts or {})
                 if not ok then break end -- XXX: continue here and return first
                                          -- non-match instead of size?
                 table.insert(batches[field_name].keys, key)
@@ -805,28 +815,35 @@ fetch_first_same = function(open_set, opts)
     end
 
     -- don't flood cache with single-key (non-batch) select results
-    if size == 1 then
-        return nil, 1
+    if size <= 1 then
+        return nil, size
     end
 
     local fetch_id = accessor:cache_fetch(batches, qcontext)
     return fetch_id, size
 end
 
--- XXX: is the function redundant because of fetch_first_same?
 fetch_resolve_list = function(prepared_object_list, opts)
+    local func_name = 'bfs_executor.fetch_resolve_list'
     local opts = opts or {}
     local accessor = opts.accessor
     local qcontext = opts.qcontext
+    local max_batch_size = opts.max_batch_size
+    local start_from = opts.start_from or 1
 
-    if not accessor:cache_is_supported() then return nil end
+    if not accessor:cache_is_supported() then return nil, 0 end
+
+    local size = 0
 
     local batches = {}
-    for i, prepared_object in ipairs(prepared_object_list) do
+    for i = 1, #prepared_object_list - start_from + 1 do
+        if i > max_batch_size then break end
+        local prepared_object = prepared_object_list[i + start_from - 1]
+
         for field_name, field_info in pairs(prepared_object.fields_info) do
             local prepared_resolve = field_info.prepared_resolve
             if prepared_resolve.is_calculated then
-                -- don't update size, because we don't add request to 'batches'
+                size = i
                 break
             end
             local prepared_select = prepared_resolve.prepared_select
@@ -837,33 +854,49 @@ fetch_resolve_list = function(prepared_object_list, opts)
             local iterator_opts = request_opts.iterator_opts
 
             if i == 1 then
-                assert(batches[field_name] == nil, 'XXX') -- XXX
+                assert(batches[field_name] == nil,
+                    ('internal error: %s: field names "%s" clash'):format(
+                    func_name, field_name))
                 batches[field_name] = {
                     collection_name = collection_name,
                     index_name = index_name,
                     keys = {key},
                     iterator_opts = iterator_opts,
                 }
+                size = i
             else
                 assert(batches[field_name].collection_name == collection_name,
-                    'XXX') -- XXX
+                    ('internal error: %s: prepared object list has ' ..
+                    'different collection names: "%s" and "%s"'):format(
+                    func_name, batches[field_name].collection_name,
+                    collection_name))
                 assert(batches[field_name].index_name == index_name,
-                    'XXX') -- XXX
-                assert(utils.is_subtable(batches[field_name].iterator_opts or
-                    {}, iterator_opts or {}), 'XXX') -- XXX
-                assert(utils.is_subtable(iterator_opts or {},
-                    batches[field_name].iterator_opts or {}), 'XXX') -- XXX
+                    ('internal error: %s: prepared object list has ' ..
+                    'different index names: "%s" and "%s"'):format(func_name,
+                    tostring(batches[field_name].index_name),
+                    tostring(index_name)))
+                local ok = utils.are_tables_same(batches[field_name].iterator_opts or
+                    {}, iterator_opts or {})
+                if not ok then -- avoid extra json.encode()
+                    assert(ok, ('internal error: %s: prepared object list ' ..
+                        'has different iterator options: "%s" and "%s"'):format(
+                        func_name,
+                        json.encode(batches[field_name].iterator_opts),
+                        json.encode(iterator_opts)))
+                end
                 table.insert(batches[field_name].keys, key)
+                size = i
             end
         end
     end
 
-    if next(batches) == nil then
-        return nil
+    -- don't flood cache with single-key (non-batch) select results
+    if size <= 1 then
+        return nil, size
     end
 
     local fetch_id = accessor:cache_fetch(batches, qcontext)
-    return fetch_id
+    return fetch_id, size
 end
 
 -- }}}
@@ -937,12 +970,14 @@ end
 ---
 --- * qcontext
 --- * accessor
+--- * max_batch_size
 ---
 --- @treturn table result of the query
 function bfs_executor.execute(schema, query_ast, variables, operation_name, opts)
     local opts = opts or {}
     local qcontext = opts.qcontext
     local accessor = opts.accessor
+    local max_batch_size = opts.max_batch_size
 
     local operation = core_query_util.getOperation(query_ast, operation_name)
     local root_object_type = schema[operation.operation]
@@ -1024,7 +1059,8 @@ function bfs_executor.execute(schema, query_ast, variables, operation_name, opts
         elseif item.prepared_object_list ~= nil then
             local child = invoke_resolve_list(item.prepared_object_list,
                 context, {qcontext = qcontext, accessor = accessor,
-                is_item_cache_only = is_item_cache_only})
+                is_item_cache_only = is_item_cache_only,
+                max_batch_size = max_batch_size})
             local child_open_set = child.open_set
             local child_cache_only_open_set = child.cache_only_open_set
             utils.expand_list(cache_only_open_set, child_cache_only_open_set)
@@ -1033,7 +1069,8 @@ function bfs_executor.execute(schema, query_ast, variables, operation_name, opts
             local open_set_to_fetch = is_item_cache_only and
                 cache_only_open_set or open_set
             local fetch_id, size = fetch_first_same(open_set_to_fetch,
-                {accessor = accessor, qcontext = qcontext})
+                {accessor = accessor, qcontext = qcontext,
+                max_batch_size = max_batch_size})
             if #open_set_to_fetch > 0 then
                 table.insert(open_set_to_fetch, math.max(2, size + 1), {
                     squash_marker = true,

graphql/impl.lua

@@ -20,6 +20,9 @@ local check = utils.check
 
 local impl = {}
 
+-- constants
+local DEF_MAX_BATCH_SIZE = 1000
+
 -- Instance of the library to provide graphql:compile() and graphql:execute()
 -- method (with creating zero configuration graphql instance under hood when
 -- calling compile() for the first time).
@@ -36,13 +39,16 @@ local default_instance
 --- @treturn table result of the operation
 local function gql_execute(qstate, variables, operation_name)
     assert(qstate.state)
+    assert(qstate.query_settings)
     local state = qstate.state
     assert(state.schema)
+    local max_batch_size = qstate.query_settings.max_batch_size or
+        state.max_batch_size
 
     check(variables, 'variables', 'table')
     check(operation_name, 'operation_name', 'string', 'nil')
+    check(max_batch_size, 'max_batch_size', 'number')
 
-    assert(qstate.query_settings)
     local qcontext = {
         query_settings = qstate.query_settings,
     }
@@ -77,6 +83,7 @@ local function gql_execute(qstate, variables, operation_name)
                 variables, operation_name, {
                     qcontext = qcontext,
                     accessor = accessor,
+                    max_batch_size = max_batch_size,
                 })
             local executor_metainfo = {
                 name = 'bfs',
@@ -131,6 +138,27 @@ local function compile_and_execute(state, query, variables, operation_name,
     return compiled_query:execute(variables, operation_name)
 end
 
+local function validate_query_settings(query_settings, opts)
+    local opts = opts or {}
+    local allow_nil = opts.allow_nil or false
+
+    local use_bfs_executor = query_settings.use_bfs_executor
+    local max_batch_size = query_settings.max_batch_size
+
+    if not allow_nil or type(use_bfs_executor) ~= 'nil' then
+        check(use_bfs_executor, 'use_bfs_executor', 'string')
+        assert(use_bfs_executor == 'never' or
+            use_bfs_executor == 'shard' or
+            use_bfs_executor == 'always',
+            "use_bfs_executor must be 'never', 'shard' (default) or " ..
+            "'always', 'got " .. tostring(use_bfs_executor))
+    end
+
+    if not allow_nil or type(max_batch_size) ~= 'nil' then
+        check(max_batch_size, 'max_batch_size', 'number')
+    end
+end
+
 --- Parse GraphQL query string, validate against the GraphQL schema and
 --- provide an object with the function to execute an operation from the
 --- request with specific variables values.
@@ -140,12 +168,13 @@ end
 --- @tparam string query text of a GraphQL query
 ---
 --- @tparam[opt] table opts the following options (described in
---- @{accessor_general.new}):
+--- @{accessor_general.new} and @{impl.new}):
 ---
 --- * resulting_object_cnt_max
 --- * fetched_object_cnt_max
 --- * timeout_ms
 --- * use_bfs_executor
+--- * max_batch_size
 ---
 --- @treturn table compiled query with `execute` and `avro_schema` functions
 local function gql_compile(state, query, opts)
@@ -168,11 +197,13 @@ local function gql_compile(state, query, opts)
             fetched_object_cnt_max = opts.fetched_object_cnt_max,
             timeout_ms = opts.timeout_ms,
             use_bfs_executor = opts.use_bfs_executor,
+            max_batch_size = opts.max_batch_size,
         }
     }
 
     accessor_general.validate_query_settings(qstate.query_settings,
         {allow_nil = true})
+    validate_query_settings(qstate.query_settings, {allow_nil = true})
 
     local gql_query = setmetatable(qstate, {
         __index = {
@@ -323,7 +354,8 @@ end
 ---     timeout_ms = <number>,
 ---     enable_mutations = <boolean>,
 ---     disable_dangling_check = <boolean>,
----     use_bfs_executor = 'never' | 'shard' (default) | 'always'
+---     use_bfs_executor = 'never' | 'shard' (default) | 'always',
+---     max_batch_size = <number>,
 --- })
 function impl.new(cfg)
     local cfg = cfg or {}
@@ -361,16 +393,12 @@ function impl.new(cfg)
 
     check(cfg.disable_dangling_check, 'disable_dangling_check', 'boolean',
         'nil')
-    check(cfg.use_bfs_executor, 'use_bfs_executor', 'string', 'nil')
-    assert(cfg.use_bfs_executor == 'never' or cfg.use_bfs_executor == 'shard'
-        or cfg.use_bfs_executor == 'always' or cfg.use_bfs_executor == nil,
-        "use_bfs_executor must be 'never', 'shard' (default) or 'always', '" ..
-        "got " .. tostring(cfg.use_bfs_executor))
-
     local state = {
         disable_dangling_check = cfg.disable_dangling_check,
         use_bfs_executor = cfg.use_bfs_executor or 'shard',
+        max_batch_size = cfg.max_batch_size or DEF_MAX_BATCH_SIZE,
     }
+    validate_query_settings(state)
     convert_schema.convert(state, cfg)
     return setmetatable(state, {
         __index = {