[couchdb] 04/16: Add fold_jobs/4 and pending_count/2, 3 to couch_jobs API

[va...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

vatamane pushed a commit to branch prototype/fdb-layer
in repository https://gitbox.apache.org/repos/asf/couchdb.git

commit d2c9dffa3aca3b3e2faed49526b0065ebb845fad
Author: Nick Vatamaniuc <va...@apache.org>
AuthorDate: Fri Aug 28 04:30:53 2020 -0400

    Add fold_jobs/4 and pending_count/2,3 to couch_jobs API
    
    fold_jobs/4 is a generalized folding API which can be used to process all the
    jobs of a particular type.
    
    get_pending_count/2,3 can be used to get the count of pending jobs. It takes
    the same options as accept, including a limit and `max_sched_time`. Just like
    accept it reads the range of pending jobs as a snapshot to avoid generating
    read conflicts.
---
 src/couch_jobs/src/couch_jobs.erl         | 37 +++++++++++++++++++++++++++----
 src/couch_jobs/src/couch_jobs_pending.erl | 32 +++++++++++++++++++++-----
 2 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/src/couch_jobs/src/couch_jobs.erl b/src/couch_jobs/src/couch_jobs.erl
index f5d6a7b..6c40f5d 100644
--- a/src/couch_jobs/src/couch_jobs.erl
+++ b/src/couch_jobs/src/couch_jobs.erl
@@ -17,10 +17,15 @@
     add/4,
     add/5,
     remove/3,
+
+    % Job monitoring
+    get_types/1,
     get_job_data/3,
     get_job_state/3,
     get_active_jobs_ids/2,
-    get_types/1,
+    fold_jobs/4,
+    pending_count/2,
+    pending_count/3,
 
     % Job processing
     accept/1,
@@ -80,6 +85,13 @@ remove(Tx, Type, JobId) when is_binary(JobId) ->
     end).
 
 
+-spec get_types(jtx()) -> [job_type()] | {error, any()}.
+get_types(Tx) ->
+    couch_jobs_fdb:tx(couch_jobs_fdb:get_jtx(Tx), fun(JTx) ->
+        couch_jobs_fdb:get_types(JTx)
+    end).
+
+
 -spec get_job_data(jtx(), job_type(), job_id()) -> {ok, job_data()} | {error,
     any()}.
 get_job_data(Tx, Type, JobId) when is_binary(JobId) ->
@@ -116,10 +128,27 @@ get_active_jobs_ids(Tx, Type) ->
     end).
 
 
--spec get_types(jtx()) -> [job_type()] | {error, any()}.
-get_types(Tx) ->
+-spec fold_jobs(jtx(), job_type(), fun(), any()) -> any().
+fold_jobs(Tx, Type, Fun, UserAcc) when is_function(Fun, 5) ->
     couch_jobs_fdb:tx(couch_jobs_fdb:get_jtx(Tx), fun(JTx) ->
-        couch_jobs_fdb:get_types(JTx)
+        maps:fold(fun(JobId, {_Seq, JobState, DataEnc}, Acc) ->
+            Data = couch_jobs_fdb:decode_data(DataEnc),
+            Fun(JTx, JobId, JobState, Data, Acc)
+        end, UserAcc, couch_jobs_fdb:get_jobs(JTx, Type))
+    end).
+
+
+-spec pending_count(jtx(), job_type()) -> integer().
+pending_count(Tx, Type) ->
+    pending_count(Tx, Type, #{}).
+
+
+-spec pending_count(jtx(), job_type(), #{}) -> integer().
+pending_count(Tx, Type, Opts) ->
+    MaxSTime = maps:get(max_sched_time, Opts, ?UNDEFINED_MAX_SCHEDULED_TIME),
+    Limit = maps:get(limit, Opts, 1024),
+    couch_jobs_fdb:tx(couch_jobs_fdb:get_jtx(Tx), fun(JTx) ->
+        couch_jobs_pending:pending_count(JTx, Type, MaxSTime, Limit)
     end).
 
 
diff --git a/src/couch_jobs/src/couch_jobs_pending.erl b/src/couch_jobs/src/couch_jobs_pending.erl
index ab53c59..a85f2fc 100644
--- a/src/couch_jobs/src/couch_jobs_pending.erl
+++ b/src/couch_jobs/src/couch_jobs_pending.erl
@@ -16,7 +16,8 @@
 -export([
     enqueue/4,
     dequeue/4,
-    remove/4
+    remove/4,
+    pending_count/4
 ]).
 
 
@@ -47,16 +48,14 @@ dequeue(#{jtx := true} = JTx, Type, _, true) ->
             {ok, JobId}
     end;
 
-dequeue(#{jtx := true} = JTx, Type, MaxPriority, _) ->
+dequeue(#{jtx := true} = JTx, Type, MaxSTime, _) ->
     #{tx := Tx, jobs_path := Jobs} = JTx,
-    Prefix = erlfdb_tuple:pack({?PENDING, Type}, Jobs),
-    StartKeySel = erlfdb_key:first_greater_than(Prefix),
-    End = erlfdb_tuple:pack({MaxPriority, <<16#FF>>}, Prefix),
-    EndKeySel = erlfdb_key:first_greater_or_equal(End),
+    {StartKeySel, EndKeySel} = get_range_selectors(JTx, Type, MaxSTime),
     case clear_random_key_from_range(Tx, StartKeySel, EndKeySel) of
         {error, not_found} ->
             {not_found, get_pending_watch(JTx, Type)};
         {ok, PendingKey} ->
+            Prefix = erlfdb_tuple:pack({?PENDING, Type}, Jobs),
             {_, JobId} = erlfdb_tuple:unpack(PendingKey, Prefix),
             {ok, JobId}
     end.
@@ -68,8 +67,29 @@ remove(#{jtx := true} = JTx, Type, JobId, STime) ->
     erlfdb:clear(Tx, Key).
 
 
+pending_count(#{jtx := true} = JTx, Type, MaxSTime, Limit) ->
+    #{tx := Tx} = JTx,
+    Opts = [
+        {limit, Limit},
+        {snapshot, true},
+        {streaming_mode, want_all}
+    ],
+    {StartSel, EndSel} = get_range_selectors(JTx, Type, MaxSTime),
+    FoldFun = fun(_Row, Cnt) -> Cnt + 1 end,
+    erlfdb:fold_range(Tx, StartSel, EndSel, FoldFun, 0, Opts).
+
+
 %% Private functions
 
+% Get pending key selectors, taking into account max scheduled time value.
+get_range_selectors(#{jtx := true} = JTx, Type, MaxSTime) ->
+    #{jobs_path := Jobs} = JTx,
+    Prefix = erlfdb_tuple:pack({?PENDING, Type}, Jobs),
+    StartKeySel = erlfdb_key:first_greater_than(Prefix),
+    End = erlfdb_tuple:pack({MaxSTime, <<16#FF>>}, Prefix),
+    EndKeySel = erlfdb_key:first_greater_or_equal(End),
+    {StartKeySel, EndKeySel}.
+
 
 % Pick a random item from the range without reading the keys in first. However
 % the constraint it that IDs should looks like random UUIDs