From e10baa80974e29de82be86d42ef8a31ca0b50652 Mon Sep 17 00:00:00 2001
From: Eric <5089238+emizzle@users.noreply.github.com>
Date: Tue, 10 Jun 2025 11:04:23 -0700
Subject: [PATCH] feat(sales): support catastrophic data loss

On startup, when a SP has filled slots, go to the SaleDownloading state instead of the SaleFilled state. This allows a SP to download slot data in case it has experienced a catastrophic loss event.
---
 codex/sales/states/downloading.nim           | 10 +-
 codex/sales/states/unknown.nim               |  6 +-
 tests/codex/sales/states/testdownloading.nim | 98 +++++++++++++++++++-
 tests/codex/sales/states/testunknown.nim     |  3 +-
 4 files changed, 113 insertions(+), 4 deletions(-)

diff --git a/codex/sales/states/downloading.nim b/codex/sales/states/downloading.nim
index 0c39b0a5..08cc1bba 100644
--- a/codex/sales/states/downloading.nim
+++ b/codex/sales/states/downloading.nim
@@ -38,6 +38,7 @@ method run*(
   let agent = SalesAgent(machine)
   let data = agent.data
   let context = agent.context
+  let market = context.market
   let reservations = context.reservations
 
   without onStore =? context.onStore:
@@ -75,8 +76,15 @@ method run*(
     trace "Starting download"
     if err =? (await onStore(request, data.slotIndex, onBlocks, isRepairing)).errorOption:
       return some State(SaleErrored(error: err, reprocessSlot: false))
-
     trace "Download complete"
+
+    # During startup, already-filled slots will download data if needed (eg
+    # catastrophic failure). If slot was already filled by current SP, go to
+    # SaleFilled.
+    let state = await market.slotState(slotId)
+    if state == SlotState.Filled:
+      return some State(SaleFilled())
+
     return some State(SaleInitialProving())
   except CancelledError as e:
     trace "SaleDownloading.run was cancelled", error = e.msgDetail
diff --git a/codex/sales/states/unknown.nim b/codex/sales/states/unknown.nim
index b714a4b9..91078ac7 100644
--- a/codex/sales/states/unknown.nim
+++ b/codex/sales/states/unknown.nim
@@ -9,6 +9,7 @@ import ./errored
 import ./proving
 import ./cancelled
 import ./payout
+import ./downloading
 
 logScope:
   topics = "marketplace sales unknown"
@@ -52,7 +53,10 @@ method run*(
         newException(UnexpectedSlotError, "Slot state on chain should not be 'free'")
       return some State(SaleErrored(error: error))
     of SlotState.Filled:
-      return some State(SaleFilled())
+      # To support catastrophic data loss, when starting, filled slots will have
+      # their data downloaded only if needed, then can proceed to the filled
+      # state.
+      return some State(SaleDownloading())
     of SlotState.Finished:
       return some State(SalePayout())
     of SlotState.Paid:
diff --git a/tests/codex/sales/states/testdownloading.nim b/tests/codex/sales/states/testdownloading.nim
index 71376fc8..34b5f99b 100644
--- a/tests/codex/sales/states/testdownloading.nim
+++ b/tests/codex/sales/states/testdownloading.nim
@@ -1,21 +1,83 @@
-import pkg/unittest2
+import std/random
+import std/times
 import pkg/questionable
 import pkg/codex/contracts/requests
 import pkg/codex/sales/states/cancelled
 import pkg/codex/sales/states/downloading
 import pkg/codex/sales/states/failed
 import pkg/codex/sales/states/filled
+import pkg/codex/sales/states/initialproving
+import pkg/codex/sales/states/errored
+import pkg/codex/sales/salesagent
+import pkg/codex/sales/salescontext
+import pkg/codex/stores/repostore
+import pkg/datastore
 import ../../examples
 import ../../helpers
+import ../../helpers/mockmarket
+import ../../../asynctest
 
 suite "sales state 'downloading'":
   let request = StorageRequest.example
   let slotIndex = request.ask.slots div 2
+  let slotId = slotId(request.id, slotIndex)
+  var market: MockMarket
+  var context: SalesContext
+  var agent: SalesAgent
   var state: SaleDownloading
+  var repo: RepoStore
+  var repoDs: Datastore
+  var metaDs: Datastore
+  var reservations: Reservations
+  let repoTmp = TempLevelDb.new()
+  let metaTmp = TempLevelDb.new()
+
+  proc createAvailability(enabled = true, until = 0.SecondsSince1970): Availability =
+    let collateralPerByte = uint8.example.u256
+    let example = Availability.example(collateralPerByte)
+    let totalSize = rand(100000 .. 200000).uint64
+    let totalCollateral = totalSize.u256 * collateralPerByte
+    let availability = waitFor reservations.createAvailability(
+      totalSize, example.duration, example.minPricePerBytePerSecond, totalCollateral,
+      enabled, until,
+    )
+    return availability.get
+
+  proc createReservation(availability: Availability): Reservation =
+    let size = rand(1 ..< availability.freeSize.int)
+    let validUntil = getTime().toUnix() + 30.SecondsSince1970
+    let reservation = waitFor reservations.createReservation(
+      availability.id, size.uint64, RequestId.example, uint64.example, 1.u256,
+      validUntil,
+    )
+    return reservation.get
 
   setup:
+    market = MockMarket.new()
+    context = SalesContext(market: market)
+    agent = newSalesAgent(context, request.id, slotIndex, request.some)
+
+    let onStore: OnStore = proc(
+        request: StorageRequest, slot: uint64, blocksCb: BlocksCb, isRepairing: bool
+    ): Future[?!void] {.gcsafe, async: (raises: [CancelledError]).} =
+      return success()
+
+    repoDs = repoTmp.newDb()
+    metaDs = metaTmp.newDb()
+    repo = RepoStore.new(repoDs, metaDs)
+    reservations = Reservations.new(repo)
+
+    let availability = createAvailability()
+    let reservation = createReservation(availability)
+
+    context.onStore = some onStore
+    agent.data.reservation = some reservation
     state = SaleDownloading.new()
 
+  teardown:
+    await repoTmp.destroyDb()
+    await metaTmp.destroyDb()
+
   test "switches to cancelled state when request expires":
     let next = state.onCancelled(request)
     check !next of SaleCancelled
@@ -27,3 +89,37 @@ suite "sales state 'downloading'":
   test "switches to filled state when slot is filled":
     let next = state.onSlotFilled(request.id, slotIndex)
     check !next of SaleFilled
+
+  test "switches to filled state after download when slot is filled":
+    market.slotState[slotId] = SlotState.Filled
+    let next = await state.run(agent)
+    check !next of SaleFilled
+
+  test "switches to initial proving state after download when slot is not filled":
+    market.slotState[slotId] = SlotState.Free
+    let next = await state.run(agent)
+    check !next of SaleInitialProving
+
+  test "calls onStore during download":
+    var onStoreCalled = false
+    let onStore: OnStore = proc(
+        request: StorageRequest, slot: uint64, blocksCb: BlocksCb, isRepairing: bool
+    ): Future[?!void] {.gcsafe, async: (raises: [CancelledError]).} =
+      onStoreCalled = true
+      return success()
+
+    context.onStore = some onStore
+    discard await state.run(agent)
+    check onStoreCalled
+
+  test "switches to error state if onStore fails":
+    var onStoreCalled = false
+    let onStore: OnStore = proc(
+        request: StorageRequest, slot: uint64, blocksCb: BlocksCb, isRepairing: bool
+    ): Future[?!void] {.gcsafe, async: (raises: [CancelledError]).} =
+      return failure "some error"
+
+    context.onStore = some onStore
+    let next = await state.run(agent)
+    check !next of SaleErrored
+    check SaleErrored(!next).error.msg == "some error"
diff --git a/tests/codex/sales/states/testunknown.nim b/tests/codex/sales/states/testunknown.nim
index 4806122f..933ee5c4 100644
--- a/tests/codex/sales/states/testunknown.nim
+++ b/tests/codex/sales/states/testunknown.nim
@@ -8,6 +8,7 @@ import pkg/codex/sales/states/filled
 import pkg/codex/sales/states/finished
 import pkg/codex/sales/states/failed
 import pkg/codex/sales/states/payout
+import pkg/codex/sales/states/downloading
 
 import ../../../asynctest
 import ../../helpers/mockmarket
@@ -49,7 +50,7 @@ suite "sales state 'unknown'":
   test "switches to filled state when on chain state is 'filled'":
     market.slotState[slotId] = SlotState.Filled
     let next = await state.run(agent)
-    check !next of SaleFilled
+    check !next of SaleDownloading
 
   test "switches to payout state when on chain state is 'finished'":
     market.slotState[slotId] = SlotState.Finished