dex: restore close-on-fill, safely

crates/core/component/dex/src/component/position_manager.rs

@@ -203,9 +203,24 @@ pub trait PositionManager: StateWrite + PositionRead {
 
     /// Record execution against an opened position.
     #[tracing::instrument(level = "debug", skip_all)]
-    async fn position_execution(&mut self, post_execution_state: position::Position) -> Result<()> {
-        self.record_proto(event::position_execution(&post_execution_state));
-        self.update_position(post_execution_state).await?;
+    async fn position_execution(&mut self, mut position: Position) -> Result<()> {
+        // Handle "close-on-fill": automatically flip the position state to "closed" if
+        // either of the reserves are zero.
+        if position.close_on_fill {
+            if position.reserves.r1 == 0u64.into() || position.reserves.r2 == 0u64.into() {
+                tracing::debug!(
+                    id = ?position.id(),
+                    r1 = ?position.reserves.r1,
+                    r2 = ?position.reserves.r2,
+                    "marking position as closed due to close-on-fill"
+                );
+                position.state = position::State::Closed;
+            }
+        }
+
+        self.record_proto(event::position_execution(&position));
+        self.update_position(position).await?;
+
         Ok(())
     }
 

crates/core/component/dex/src/component/tests.rs

@@ -9,6 +9,7 @@ use penumbra_num::Amount;
 use rand_core::OsRng;
 
 use crate::lp::action::PositionOpen;
+use crate::lp::{position, SellOrder};
 use crate::DexParameters;
 use crate::{
     component::{
@@ -257,6 +258,59 @@ async fn single_limit_order() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+/// Builds a simple order book with a two orders, fills against them both,
+/// and checks that one of the orders is auto-closed.
+async fn check_close_on_fill() -> anyhow::Result<()> {
+    let storage = TempStorage::new().await?.apply_minimal_genesis().await?;
+    let mut state = Arc::new(StateDelta::new(storage.latest_snapshot()));
+    let mut state_tx = state.try_begin_transaction().unwrap();
+
+    let gm = asset::Cache::with_known_assets().get_unit("gm").unwrap();
+
+    let mut position_1 = SellOrder::parse_str("100gm@1gn")?.into_position(OsRng);
+    position_1.close_on_fill = true;
+    let position_2 = SellOrder::parse_str("[email protected]")?.into_position(OsRng);
+
+    let position_1_id = position_1.id();
+    let position_2_id = position_2.id();
+
+    state_tx.open_position(position_1.clone()).await.unwrap();
+    state_tx.open_position(position_2.clone()).await.unwrap();
+
+    // Now we have the following liquidity:
+    //
+    // 100gm@1gn (auto-closing)
+    // [email protected]
+    //
+    // We therefore expect that trading 100gn + 110gn will exhaust both positions.
+    // Attempting to trade a bit more than that ensures we completely fill both,
+    // without worrying about rounding.
+    // Because we're just testing the DEX internals, we need to trigger fill_route manually.
+    let input = "220gn".parse::<Value>().unwrap();
+    let route = [gm.id()];
+    let execution = FillRoute::fill_route(&mut state_tx, input, &route, None).await?;
+
+    let unfilled = input.amount.checked_sub(&execution.input.amount).unwrap();
+
+    // Check that we got the execution we expected.
+    assert_eq!(unfilled, "10gn".parse::<Value>().unwrap().amount);
+    assert_eq!(execution.output, "200gm".parse::<Value>().unwrap());
+
+    // Now grab both position states:
+    let position_1_post_exec = state_tx.position_by_id(&position_1_id).await?.unwrap();
+    let position_2_post_exec = state_tx.position_by_id(&position_2_id).await?.unwrap();
+
+    dbg!(&position_1_post_exec);
+    dbg!(&position_2_post_exec);
+
+    // Check that position 1 was auto-closed but position 2 wasn't:
+    assert_eq!(position_1_post_exec.state, position::State::Closed);
+    assert_eq!(position_2_post_exec.state, position::State::Opened);
+
+    Ok(())
+}
+
 #[tokio::test]
 /// Try to execute against multiple positions, mainly testing that the order-book traversal
 /// is done correctly.

crates/core/component/dex/src/lp/position.rs

@@ -33,7 +33,7 @@ pub struct Position {
     /// sequence of stateful NFTs based on the [`Id`].
     pub nonce: [u8; 32],
     /// Set to `true` if a position is a limit-order, meaning that it will be closed after being
-    /// filled against. Note that this is not currently supported in the dex state machine.
+    /// filled against.
     pub close_on_fill: bool,
 }