fix: list invocations pagination

seriousben · seriousben · commit 827edea8526e · 2025-04-14T11:24:35.000-04:00
diff --git a/server/src/http_objects.rs b/server/src/http_objects.rs
@@ -421,7 +421,8 @@ pub struct CreateNamespaceResponse {
 #[derive(Debug, Serialize, Deserialize, ToSchema)]
 pub struct GraphInvocations {
     pub invocations: Vec<Invocation>,
-    pub cursor: Option<String>,
+    pub prev_cursor: Option<String>,
+    pub next_cursor: Option<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize, ToSchema)]
diff --git a/server/src/routes.rs b/server/src/routes.rs
@@ -636,7 +636,7 @@ async fn graph_invocations(
         Some(CursorDirection::Backward) => Some(state_store::scanner::CursorDirection::Backward),
         None => None,
     };
-    let (invocation_ctxs, cursor) = state
+    let (invocation_ctxs, prev_cursor, next_cursor) = state
         .indexify_state
         .reader()
         .list_invocations(
@@ -651,11 +651,13 @@ async fn graph_invocations(
     for invocation_ctx in invocation_ctxs {
         invocations.push(invocation_ctx.into());
     }
-    let cursor = cursor.map(|c| BASE64_STANDARD.encode(c));
+    let prev_cursor = prev_cursor.map(|c| BASE64_STANDARD.encode(c));
+    let next_cursor = next_cursor.map(|c| BASE64_STANDARD.encode(c));
 
     Ok(Json(GraphInvocations {
         invocations,
-        cursor,
+        prev_cursor,
+        next_cursor,
     }))
 }
 
@@ -971,7 +973,7 @@ async fn find_invocation(
     Ok(Json(invocation_ctx.into()))
 }
 
-/// Delete a specific invocation  
+/// Delete a specific invocation
 #[utoipa::path(
     delete,
     path = "/namespaces/{namespace}/compute_graphs/{compute_graph}/invocations/{invocation_id}",
diff --git a/server/state_store/src/scanner.rs b/server/state_store/src/scanner.rs
@@ -17,6 +17,7 @@ use data_model::{
     TaskOutputsIngestedEvent,
     UnprocessedStateChanges,
 };
+use indexify_utils::get_epoch_time_in_ms;
 use metrics::Timer;
 use opentelemetry::KeyValue;
 use rocksdb::{Direction, IteratorMode, ReadOptions, TransactionDB};
@@ -493,21 +494,21 @@ impl StateReader {
         cursor: Option<&[u8]>,
         limit: usize,
         direction: Option<CursorDirection>,
-    ) -> Result<(Vec<GraphInvocationCtx>, Option<Vec<u8>>)> {
+    ) -> Result<(Vec<GraphInvocationCtx>, Option<Vec<u8>>, Option<Vec<u8>>)> {
         let kvs = &[KeyValue::new("op", "list_invocations")];
         let _timer = Timer::start_with_labels(&self.metrics.state_read, kvs);
-        let key_prefix = [namespace.as_bytes(), b"|", compute_graph.as_bytes()].concat();
+        let key_prefix = [namespace.as_bytes(), b"|", compute_graph.as_bytes(), b"|"].concat();
 
         let direction = direction.unwrap_or(CursorDirection::Forward);
         let mut read_options = ReadOptions::default();
         read_options.set_readahead_size(10_194_304);
 
         let mut upper_bound = key_prefix.clone();
-        upper_bound.push(0xff);
+        upper_bound.extend(&get_epoch_time_in_ms().to_be_bytes());
         read_options.set_iterate_upper_bound(upper_bound);
 
         let mut lower_bound = key_prefix.clone();
-        lower_bound.push(0x00);
+        lower_bound.extend(&(0 as u64).to_be_bytes());
         read_options.set_iterate_lower_bound(lower_bound);
 
         let mut iter = self.db.raw_iterator_cf_opt(
@@ -516,29 +517,67 @@ impl StateReader {
         );
 
         match cursor {
-            Some(cursor) => iter.seek(cursor),
-            None => iter.seek_to_last(),
+            Some(cursor) => {
+                match direction {
+                    CursorDirection::Backward => iter.seek(cursor),
+                    CursorDirection::Forward => iter.seek_for_prev(cursor),
+                }
+                // Skip the first item (cursor position)
+                if iter.valid() {
+                    match direction {
+                        CursorDirection::Forward => iter.prev(),
+                        CursorDirection::Backward => iter.next(),
+                    }
+                }
+            }
+            None => match direction {
+                CursorDirection::Backward => iter.seek_to_first(), // Start at beginning of range
+                CursorDirection::Forward => iter.seek_to_last(),   // Start at end of range
+            },
         }
 
         let mut rows = Vec::new();
         let mut next_cursor = None;
+        let mut prev_cursor = None;
 
         // Collect results
-        while iter.valid() {
+        while iter.valid() && rows.len() < limit {
             if let Some((key, _v)) = iter.item() {
-                if rows.len() < limit {
-                    rows.push(key.to_vec());
-                } else {
-                    next_cursor = Some(key.to_vec());
-                    break;
-                }
+                rows.push(key.to_vec());
             } else {
-                break;
+                break; // No valid item found
             }
 
+            // Move the iterator after capturing the current item
             match direction {
-                CursorDirection::Backward => iter.next(),
                 CursorDirection::Forward => iter.prev(),
+                CursorDirection::Backward => iter.next(),
+            }
+        }
+
+        // Check if there are more items after our limit
+        if iter.valid() {
+            let key = rows.last().cloned();
+            match direction {
+                CursorDirection::Forward => {
+                    next_cursor = key;
+                }
+                CursorDirection::Backward => {
+                    prev_cursor = key;
+                }
+            }
+        }
+
+        // Set the previous cursor if we have a valid item
+        if cursor.is_some() {
+            let key = rows.first().cloned();
+            match direction {
+                CursorDirection::Forward => {
+                    prev_cursor = key;
+                }
+                CursorDirection::Backward => {
+                    next_cursor = key;
+                }
             }
         }
 
@@ -556,13 +595,22 @@ impl StateReader {
             }
         }
 
+        match direction {
+            CursorDirection::Forward => {}
+            CursorDirection::Backward => {
+                // We keep the ordering the same even if we traverse in the opposite direction
+                invocation_prefixes.reverse();
+            }
+        }
+
         let invocations = self.get_rows_from_cf_multi_key::<GraphInvocationCtx>(
             invocation_prefixes.iter().map(|v| v.as_slice()).collect(),
             IndexifyObjectsColumns::GraphInvocationCtx,
         )?;
 
         Ok((
             invocations.into_iter().filter_map(|v| v).collect(),
+            prev_cursor,
             next_cursor,
         ))
     }
diff --git a/server/ui/src/routes/Namespace/IndividualComputeGraphPage.tsx b/server/ui/src/routes/Namespace/IndividualComputeGraphPage.tsx
@@ -1,4 +1,12 @@
-import { Box, Breadcrumbs, Typography, Stack, Chip, Button, CircularProgress } from '@mui/material'
+import {
+  Box,
+  Breadcrumbs,
+  Typography,
+  Stack,
+  Chip,
+  Button,
+  CircularProgress,
+} from '@mui/material'
 import { TableDocument } from 'iconsax-react'
 import NavigateNextIcon from '@mui/icons-material/NavigateNext'
 import NavigateBeforeIcon from '@mui/icons-material/NavigateBefore'
@@ -14,13 +22,18 @@ import axios from 'axios'
 import { getIndexifyServiceURL } from '../../utils/helpers'
 
 const IndividualComputeGraphPage = () => {
-  const { invocationsList, computeGraph, namespace, cursor } = useLoaderData() as IndividualComputeGraphLoaderData
+  const {
+    invocationsList,
+    computeGraph,
+    namespace,
+    prevCursor: prevCursorLoader,
+    nextCursor: nextCursorLoader,
+  } = useLoaderData() as IndividualComputeGraphLoaderData
 
   const [invocations, setInvocations] = useState<Invocation[]>(invocationsList)
   const [isLoading, setIsLoading] = useState(false)
-  const [currentCursor, setCurrentCursor] = useState<string | null>(null)
-  const [nextCursor, setNextCursor] = useState<string | null>(cursor)
-  const [cursorHistory, setCursorHistory] = useState<string[]>([])
+  const [prevCursor, setPrevCursor] = useState<string | null>(prevCursorLoader)
+  const [nextCursor, setNextCursor] = useState<string | null>(nextCursorLoader)
 
   const handleDelete = useCallback((updatedList: Invocation[]) => {
     const sortedList = [...updatedList].sort(
@@ -29,56 +42,49 @@ const IndividualComputeGraphPage = () => {
     setInvocations(sortedList)
   }, [])
 
-  const fetchInvocations = useCallback(async (cursor: string | null, direction: 'forward' | 'backward') => {
-    setIsLoading(true)
-    try {
-      const serviceURL = getIndexifyServiceURL()
-      const limit = 20
-      const url = `${serviceURL}/namespaces/${namespace}/compute_graphs/${computeGraph.name}/invocations?limit=${limit}${
-        cursor ? `&cursor=${cursor}` : ''
-      }&direction=${direction}`
-      
-      const response = await axios.get(url)
-      const data = response.data
-      
-      setInvocations([...data.invocations].sort(
-        (a, b) => (b.created_at ?? 0) - (a.created_at ?? 0)
-      ))
-      
-      if (direction === 'forward') {
-        if (cursor) {
-          setCursorHistory(prev => [...prev, cursor])
-        }
-        setCurrentCursor(cursor)
-        setNextCursor(data.cursor || null)
-      } else {
-        if (cursorHistory.length > 0) {
-          setCursorHistory(prev => prev.slice(0, -1))
-        }
-        setCurrentCursor(cursorHistory.length > 1 ? cursorHistory[cursorHistory.length - 2] : null)
-        setNextCursor(data.cursor || null)
+  const fetchInvocations = useCallback(
+    async (cursor: string | null, direction: 'forward' | 'backward') => {
+      setIsLoading(true)
+      try {
+        const serviceURL = getIndexifyServiceURL()
+        const limit = 20
+        const url = `${serviceURL}/namespaces/${namespace}/compute_graphs/${
+          computeGraph.name
+        }/invocations?limit=${limit}${
+          cursor ? `&cursor=${cursor}` : ''
+        }&direction=${direction}`
+
+        const response = await axios.get(url)
+        const data = response.data
+
+        setInvocations([...data.invocations])
+
+        setPrevCursor(data.prev_cursor)
+        setNextCursor(data.next_cursor)
+        console.log(direction, {
+          prevCursor: data.prev_cursor,
+          nextCursor: data.next_cursor,
+        })
+      } catch (error) {
+        console.error('Error fetching invocations:', error)
+      } finally {
+        setIsLoading(false)
       }
-    } catch (error) {
-      console.error("Error fetching invocations:", error)
-    } finally {
-      setIsLoading(false)
-    }
-  }, [namespace, computeGraph.name, cursorHistory])
+    },
+    [namespace, computeGraph.name]
+  )
 
   const handleNextPage = useCallback(() => {
-    const cursor = nextCursor || currentCursor
-    if (cursor) {
-      fetchInvocations(cursor, 'forward')
+    if (nextCursor) {
+      fetchInvocations(nextCursor, 'forward')
     }
-  }, [nextCursor, currentCursor, fetchInvocations])
+  }, [nextCursor, fetchInvocations])
 
   const handlePreviousPage = useCallback(() => {
-    if (cursorHistory.length > 0) {
-      const prevCursor = cursorHistory[cursorHistory.length - 1]
+    if (prevCursor) {
       fetchInvocations(prevCursor, 'backward')
     }
-  }, [cursorHistory, fetchInvocations])
-
+  }, [prevCursor, fetchInvocations])
   return (
     <Stack direction="column" spacing={3}>
       <Breadcrumbs
@@ -132,29 +138,29 @@ const IndividualComputeGraphPage = () => {
           onDelete={handleDelete}
         />
 
-        <Box 
-          sx={{ 
-            display: 'flex', 
-            justifyContent: 'space-between', 
+        <Box
+          sx={{
+            display: 'flex',
+            justifyContent: 'space-between',
             mt: 2,
-            alignItems: 'center'
+            alignItems: 'center',
           }}
         >
           <Button
             startIcon={<NavigateBeforeIcon />}
             onClick={handlePreviousPage}
-            disabled={cursorHistory.length === 0 || isLoading}
+            disabled={!prevCursor || isLoading}
             variant="outlined"
           >
             Previous
           </Button>
-          
+
           {isLoading && <CircularProgress size={24} />}
-          
+
           <Button
             endIcon={<NavigateNextIcon />}
             onClick={handleNextPage}
-            disabled={isLoading}
+            disabled={!nextCursor || isLoading}
             variant="outlined"
           >
             Next
diff --git a/server/ui/src/routes/Namespace/types.ts b/server/ui/src/routes/Namespace/types.ts
@@ -19,7 +19,8 @@ export interface ComputeGraphLoaderData extends NamespaceLoaderData {
 export interface IndividualComputeGraphLoaderData extends NamespaceLoaderData {
   invocationsList: Invocation[]
   computeGraph: ComputeGraph
-  cursor: string | null;
+  prevCursor: string | null
+  nextCursor: string | null
   direction?: string
 }
 
diff --git a/server/ui/src/utils/loaders.ts b/server/ui/src/utils/loaders.ts

Original file line number	Diff line number	Diff line change
`@@ -421,7 +421,8 @@ pub struct CreateNamespaceResponse {`
`421`	`421`	`#[derive(Debug, Serialize, Deserialize, ToSchema)]`
`422`	`422`	`pub struct GraphInvocations {`
`423`	`423`	`pub invocations: Vec<Invocation>,`
`424`		`- pub cursor: Option<String>,`
	`424`	`+ pub prev_cursor: Option<String>,`
	`425`	`+ pub next_cursor: Option<String>,`
`425`	`426`	`}`
`426`	`427`
`427`	`428`	`#[derive(Debug, Serialize, Deserialize, ToSchema)]`
Original file line number	Diff line number	Diff line change
`@@ -19,7 +19,8 @@ export interface ComputeGraphLoaderData extends NamespaceLoaderData {`
`19`	`19`	`export interface IndividualComputeGraphLoaderData extends NamespaceLoaderData {`
`20`	`20`	`invocationsList: Invocation[]`
`21`	`21`	`computeGraph: ComputeGraph`
`22`		`- cursor: string \| null;`
	`22`	`+ prevCursor: string \| null`
	`23`	`+ nextCursor: string \| null`
`23`	`24`	`direction?: string`
`24`	`25`	`}`
`25`	`26`