fix: stabilize process uploads and scheduling

src/core/http/hb_client_remote.erl

@@ -112,8 +112,11 @@ upload(Msg, Opts, _CommitmentDevice) ->
     hb_ao:raw(
         <<"arweave@2.9">>,
         <<"tx">>,
-        #{},
-        Msg#{ <<"method">> => <<"POST">> },
+        Msg,
+        #{
+            <<"method">> => <<"POST">>,
+            <<"target">> => <<"base">>
+        },
         Opts
     ).
 
@@ -154,3 +157,22 @@ upload_single_layer_message_test() ->
     Result = upload(Committed, Opts, <<"ans104@1.0">>),
     ?event({upload_result, Result}),
     ?assertMatch({ok, _}, Result).
+
+%% @doc A message's AO target is transaction data, not the arweave device's
+%% request/base selector. Upload it as the base so a process ID target is not
+%% misinterpreted as a selector path.
+upload_targeted_message_test() ->
+    Opts = upload_test_opts(),
+    Msg = #{
+        <<"target">> => hb_util:human_id(<<0:256>>),
+        <<"data">> => <<"TEST">>
+    },
+    Committed =
+        hb_message:commit(
+            Msg,
+            Opts,
+            <<"ans104@1.0">>
+        ),
+    Result = upload(Committed, Opts, <<"ans104@1.0">>),
+    ?event({upload_result, Result}),
+    ?assertMatch({ok, _}, Result).

src/core/store/hb_store_rocksdb.erl

@@ -224,14 +224,9 @@ link_path(_, Key1, Key1) ->
 link_path(Opts, Existing, New) ->
     ExistingBin = convert_if_list(Existing),
     NewBin = convert_if_list(New),
-
-    % Create: NewValue -> ExistingBin
-    case do_read(Opts, NewBin) of
-        not_found ->
-            do_write(Opts, NewBin, encode_value(link, ExistingBin));
-        _ ->
-            ok
-    end.
+    % Create or replace: NewValue -> ExistingBin. Replacement matches the
+    % semantics of the FS, LMDB, and volatile stores.
+    do_write(Opts, NewBin, encode_value(link, ExistingBin)).
 
 %% @doc List all items registered in rocksdb store. Should be used only
 %% for testing/debugging, as the underlying operation is doing full traversal
@@ -457,8 +452,13 @@ write_read_test_() ->
                 ),
                 ok = link(#{}, #{ <<"test_key">> => <<"test_key2">> }, #{}),
                 {ok, Value} = read(#{}, #{ <<"read">> => <<"test_key">> }, #{}),
-
-                ?assertEqual(<<"value_under_linked_key">>, Value)
+                ?assertEqual(<<"value_under_linked_key">>, Value),
+                ok = write(#{}, #{ <<"test_key3">> => <<"replacement">> }, #{}),
+                ok = link(#{}, #{ <<"test_key">> => <<"test_key3">> }, #{}),
+                ?assertEqual(
+                    {ok, <<"replacement">>},
+                    read(#{}, #{ <<"read">> => <<"test_key">> }, #{})
+                )
             end}
         ]}.
 

src/preloaded/process/dev_scheduler_cache.erl

@@ -1,6 +1,6 @@
 %%% @doc A module that provides a cache for scheduler assignments and locations.
 -module(dev_scheduler_cache).
--export([write/2, write_spawn/2, read/3]).
+-export([write/2, unlink/2, write_spawn/2, read/3]).
 -export([list/2, latest/2]).
 -include("include/hb.hrl").
 -include_lib("eunit/include/eunit.hrl").
@@ -41,14 +41,8 @@ write(RawAssignment, RawOpts) ->
             % slot on the process to the underlying data.
             ok = hb_store:link(
                 Store,
-                #{
-                    hb_path:to_binary([
-                        ?SCHEDULER_CACHE_PREFIX,
-                        <<"assignments">>,
-                        hb_util:human_id(ProcID),
-                        hb_ao:normalize_key(Slot)
-                    ]) => hb_message:id(Assignment, signed, Opts)
-                },
+                #{ assignment_path(ProcID, Slot) =>
+                    hb_message:id(Assignment, signed, Opts) },
                 Opts
             ),
             ok;
@@ -57,6 +51,34 @@ write(RawAssignment, RawOpts) ->
             {error, Reason}
     end.
 
+%% @doc Remove a failed assignment from the usable schedule while retaining
+%% its content-addressed data. Stores do not expose a destructive unlink API,
+%% so replace the slot link with a deliberately dangling rollback link.
+unlink(RawAssignment, RawOpts) ->
+    Assignment = hb_cache:ensure_all_loaded(RawAssignment, RawOpts),
+    Opts = opts(RawOpts),
+    Store = hb_opts:get(store, no_viable_store, Opts),
+    ProcID = hb_ao:get(<<"process">>, Assignment, Opts),
+    Slot = hb_ao:get(<<"slot">>, Assignment, Opts),
+    AssignmentID = hb_message:id(Assignment, signed, Opts),
+    ?event(
+        {unlinking_assignment,
+            {proc_id, ProcID},
+            {slot, Slot},
+            {assignment, AssignmentID}
+        }
+    ),
+    hb_store:link(
+        Store,
+        #{ assignment_path(ProcID, Slot) =>
+            hb_path:to_binary([
+                ?SCHEDULER_CACHE_PREFIX,
+                <<"rolled-back">>,
+                AssignmentID
+            ]) },
+        Opts
+    ).
+
 %% @doc Write the initial assignment message to the cache.
 write_spawn(RawInitMessage, Opts) ->
     InitMessage = hb_cache:ensure_all_loaded(RawInitMessage, Opts),
@@ -68,12 +90,7 @@ read(ProcID, Slot, Opts) when is_integer(Slot) ->
 read(ProcID, Slot, RawOpts) ->
     Opts = opts(RawOpts),
     Store = hb_opts:get(store, no_viable_store, Opts),
-    P1 = hb_path:to_binary([
-        ?SCHEDULER_CACHE_PREFIX,
-        <<"assignments">>,
-        hb_util:human_id(ProcID),
-        Slot
-    ]),
+    P1 = assignment_path(ProcID, Slot),
     ?event(
         {read_assignment,
             {proc_id, ProcID},
@@ -100,6 +117,8 @@ read(ProcID, Slot, RawOpts) ->
                             ?event({normalized_aos2_assignment, Norm}),
                             {ok, Norm};
                         <<"ao.N.1">> ->
+                            {ok, hb_cache:ensure_all_loaded(Assignment, Opts)};
+                        not_found ->
                             {ok, hb_cache:ensure_all_loaded(Assignment, Opts)}
                     end;
                 {error, not_found} ->
@@ -114,14 +133,36 @@ read(ProcID, Slot, RawOpts) ->
 %% @doc Get the assignments for a process.
 list(ProcID, RawOpts) ->
     Opts = opts(RawOpts),
-    hb_cache:list_numbered(
-        hb_path:to_binary([
-            ?SCHEDULER_CACHE_PREFIX,
-            <<"assignments">>,
-            hb_util:human_id(ProcID)
-        ]),
-        Opts
-    ).
+    Slots =
+        hb_cache:list_numbered(
+            hb_path:to_binary([
+                ?SCHEDULER_CACHE_PREFIX,
+                <<"assignments">>,
+                hb_util:human_id(ProcID)
+            ]),
+            Opts
+        ),
+    trim_unavailable_tail(ProcID, Slots, Opts).
+
+%% Rolled-back assignments can only occur at the unconfirmed tail because the
+%% scheduler does not advance its state in remote-confirmation mode. Usually
+%% this checks one slot; recursion also tolerates multiple interrupted retries.
+trim_unavailable_tail(_ProcID, [], _Opts) -> [];
+trim_unavailable_tail(ProcID, Slots, Opts) ->
+    Latest = lists:max(Slots),
+    case read(ProcID, Latest, Opts) of
+        {ok, _} -> Slots;
+        not_found ->
+            trim_unavailable_tail(ProcID, lists:delete(Latest, Slots), Opts)
+    end.
+
+assignment_path(ProcID, Slot) ->
+    hb_path:to_binary([
+        ?SCHEDULER_CACHE_PREFIX,
+        <<"assignments">>,
+        hb_util:human_id(ProcID),
+        hb_ao:normalize_key(Slot)
+    ]).
 
 %% @doc Get the latest assignment from the cache.
 latest(ProcID, RawOpts) ->
@@ -178,6 +219,20 @@ volatile_schedule_test() ->
     ?assertMatch({1, _}, latest(ProcID, Opts)),
     {ok, ReadAssignment} = read(ProcID, 1, Opts),
     ?assertEqual(ReadAssignment, hb_message:normalize_commitments(Assignment, Opts)),
+    ?assertEqual(ok, unlink(Assignment, Opts)),
+    ?assertEqual([], list(ProcID, Opts)),
+    ?assertEqual(not_found, read(ProcID, 1, Opts)),
+    ?assertEqual(ok, write(Assignment, Opts)),
+    ?assertMatch({ok, _}, read(ProcID, 1, Opts)),
+    LMDBStore = hb_test_utils:test_store(hb_store_lmdb, <<"scheduler-unlink">>),
+    LMDBOpts = #{ <<"store">> => [LMDBStore] },
+    hb_store:start(LMDBStore),
+    ?assertEqual(ok, write(Assignment, LMDBOpts)),
+    ?assertEqual(ok, unlink(Assignment, LMDBOpts)),
+    ?assertEqual([], list(ProcID, LMDBOpts)),
+    ?assertEqual(ok, write(Assignment, LMDBOpts)),
+    ?assertMatch({ok, _}, read(ProcID, 1, LMDBOpts)),
+    hb_store:reset(LMDBStore),
     hb_store:stop(VolStore),
     hb_store:reset(VolStore),
     hb_store:start(VolStore),

src/preloaded/process/dev_scheduler_server.erl

@@ -263,17 +263,7 @@ do_assign(State, Message, ReplyPID) ->
                 State
             ),
             ?event(writes_complete),
-            ?event(uploading_message),
-            hb_client_remote:upload(Message, Opts),
-            hb_client_remote:upload(Assignment, Opts),
-            ?event(uploads_complete),
-            maybe_inform_recipient(
-                remote_confirmation,
-                ReplyPID,
-                Message,
-                Assignment,
-                State
-            )
+            upload_after_write(Message, Assignment, ReplyPID, State)
         end,
     case hb_opts:get(scheduling_mode, sync, Opts) of
         aggressive ->
@@ -288,6 +278,64 @@ do_assign(State, Message, ReplyPID) ->
         base_state_hashpath := next_hashpath(BaseStateHashpath, Assignment, State)
     }.
 
+%% @doc Upload a durably written assignment. Local and aggressive confirmation
+%% have already replied by this point, so an upload failure must not roll their
+%% scheduler state back. Remote confirmation has not replied, so remove the
+%% local schedule link and re-raise to keep the previous scheduler state.
+upload_after_write(Message, Assignment, ReplyPID, State) ->
+    Opts = maps:get(opts, State),
+    try
+        ?event(uploading_message),
+        ok = ensure_upload_succeeded(hb_client_remote:upload(Message, Opts)),
+        ok = ensure_upload_succeeded(hb_client_remote:upload(Assignment, Opts)),
+        ?event(uploads_complete),
+        maybe_inform_recipient(
+            remote_confirmation,
+            ReplyPID,
+            Message,
+            Assignment,
+            State
+        )
+    catch
+        Class:Reason:Stack ->
+            ?event(error,
+                {assignment_upload_failed_after_write,
+                    {class, Class},
+                    {reason, Reason},
+                    {trace, Stack},
+                    {process, maps:get(id, State)},
+                    {slot, hb_ao:get(<<"slot">>, Assignment, Opts)}
+                }
+            ),
+            case maps:get(mode, State) of
+                aggressive -> ok;
+                local_confirmation -> ok;
+                remote_confirmation ->
+                    ok = dev_scheduler_cache:unlink(Assignment, Opts),
+                    erlang:raise(Class, Reason, Stack);
+                _ ->
+                    erlang:raise(Class, Reason, Stack)
+            end
+    end.
+
+%% @doc Turn failed upload return values into exceptions so the confirmation-
+%% mode-specific recovery in `upload_after_write/4' is always applied.
+ensure_upload_succeeded({ok, Results}) when is_list(Results) ->
+    case [Result || Result <- Results, not upload_result_succeeded(Result)] of
+        [] -> ok;
+        Failed -> erlang:error({upload_failed, Failed})
+    end;
+ensure_upload_succeeded({error, Reason}) ->
+    erlang:error({upload_failed, Reason});
+ensure_upload_succeeded({failure, Reason}) ->
+    erlang:error({upload_failed, Reason});
+ensure_upload_succeeded(Result) ->
+    erlang:error({upload_failed, Result}).
+
+upload_result_succeeded(ok) -> true;
+upload_result_succeeded({ok, _}) -> true;
+upload_result_succeeded(_) -> false.
+
 %% @doc Commit to the assignment using all of our appropriate wallets.
 commit_assignment(BaseAssignment, State) ->
     Wallets = maps:get(wallets, State),
@@ -368,7 +416,43 @@ new_proc_test() ->
         #{ current := 2 },
         dev_scheduler_server:info(dev_scheduler_registry:find(ID))
     ).
-
+
+%% @doc A returned upload error must invalidate an assignment that has not yet
+%% received remote confirmation.
+remote_confirmation_upload_failure_unlinks_test() ->
+    Wallet = hb:wallet(),
+    Store = hb_test_utils:test_store(hb_store_fs, <<"remote-upload-failure">>),
+    Opts = #{
+        <<"priv-wallet">> => Wallet,
+        <<"store">> => [Store]
+    },
+    hb_store:start(Store),
+    ProcID = hb_util:human_id(crypto:strong_rand_bytes(32)),
+    Message = hb_message:commit(
+        #{ <<"target">> => ProcID, <<"data">> => <<"test">> },
+        Opts,
+        <<"httpsig@1.0">>
+    ),
+    Assignment = hb_message:commit(
+        #{
+            <<"variant">> => <<"ao.N.1">>,
+            <<"process">> => ProcID,
+            <<"slot">> => 0,
+            <<"hash-chain">> => <<"test-hash-chain">>
+        },
+        Opts,
+        <<"httpsig@1.0">>
+    ),
+    ok = dev_scheduler_cache:write(Assignment, Opts),
+    ?assertEqual([0], dev_scheduler_cache:list(ProcID, Opts)),
+    State = #{ id => ProcID, mode => remote_confirmation, opts => Opts },
+    ?assertError(
+        {upload_failed, [{error, no_httpsig_bundler}]},
+        upload_after_write(Message, Assignment, self(), State)
+    ),
+    ?assertEqual([], dev_scheduler_cache:list(ProcID, Opts)),
+    ?assertEqual(not_found, dev_scheduler_cache:read(ProcID, 0, Opts)),
+    hb_store:reset(Store).
 
 benchmark_test() ->
     BenchTime = 1,