feat(iisa): add load balancing context with batch query optimization

bin/dipper-service/src/registry.rs

@@ -146,6 +146,20 @@ impl AgreementRegistry for RegistryProvider {
             .collect())
     }
 
+    async fn get_active_indexing_agreements_by_indexer_ids(
+        &self,
+        indexer_ids: &[IndexerId],
+    ) -> RegistryResult<Vec<IndexingAgreement>> {
+        Ok(self
+            .inner
+            .get_active_indexing_agreements_by_indexer_ids(indexer_ids)
+            .await?
+            .into_iter()
+            .map(IndexingAgreement::try_from)
+            .filter_map(Result::ok)
+            .collect())
+    }
+
     async fn get_indexing_agreements_by_indexing_request_id(
         &self,
         request_id: &IndexingRequestId,

bin/dipper-service/src/registry/agreement.rs

@@ -44,6 +44,15 @@ pub trait AgreementRegistry {
         indexer_id: &IndexerId,
     ) -> RegistryResult<Vec<IndexingAgreement>>;
 
+    /// Get all active agreements for multiple indexers in a single query.
+    ///
+    /// Returns agreements that are in `Created` or `Accepted` status for any of the
+    /// provided indexer IDs. This is more efficient than querying each indexer separately.
+    async fn get_active_indexing_agreements_by_indexer_ids(
+        &self,
+        indexer_ids: &[IndexerId],
+    ) -> RegistryResult<Vec<IndexingAgreement>>;
+
     /// Get all agreements by associated indexing request ID.
     async fn get_indexing_agreements_by_indexing_request_id(
         &self,

bin/dipper-service/src/worker/handlers.rs

@@ -2,6 +2,7 @@ mod find_indexer_for_indexing_request;
 mod process_indexing_agreement_cancellation;
 mod process_indexing_request_cancellation;
 mod process_new_indexing_request;
+mod selection_context;
 mod send_indexing_agreement_cancellation;
 mod send_indexing_agreement_proposal;
 

bin/dipper-service/src/worker/handlers/find_indexer_for_indexing_request.rs

@@ -6,6 +6,7 @@ use jsonrpsee::core::Serialize;
 use serde::Deserialize;
 use thegraph_core::{DeploymentId, alloy::primitives::ChainId};
 
+use super::selection_context::gather_selection_context;
 use crate::{
     config::{IndexingAgreementChainPrices, IndexingAgreementConfig},
     network::NetworkProvider,
@@ -97,9 +98,12 @@ where
         return Ok(());
     }
 
+    // Gather load balancing context for IISA
+    let context = gather_selection_context(&ctx.registry, deployment_id, &indexers).await?;
+
     let Some(candidate) = ctx
         .iisa
-        .select_one(*deployment_id, indexers)
+        .select_one(*deployment_id, indexers, &context)
         .await
         .map_err(|err| JobError::Fatal(err.into()))?
     else {

bin/dipper-service/src/worker/handlers/process_new_indexing_request.rs

@@ -4,6 +4,7 @@ use dipper_core::ids::IndexingRequestId;
 use dipper_iisa::{CandidateSelection, Indexer as IndexerCandidate};
 use thegraph_core::{DeploymentId, alloy::primitives::ChainId};
 
+use super::selection_context::gather_selection_context;
 use crate::{
     config::{IndexingAgreementChainPrices, IndexingAgreementConfig},
     network::NetworkProvider,
@@ -57,7 +58,7 @@ where
     W: WorkerQueue,
     I: CandidateSelection,
 {
-    // Get the indexers that are not indexing the deployment amd treat it as the raw candidate list
+    // Get the indexers that are not indexing the deployment and treat it as the raw candidate list
     // and pass it to the IISA to get the final list of candidates
     let indexers = ctx
         .network
@@ -76,9 +77,12 @@ where
         return Ok(());
     }
 
+    // Gather load balancing context for IISA
+    let context = gather_selection_context(&ctx.registry, deployment_id, &indexers).await?;
+
     let candidates = ctx
         .iisa
-        .select(*deployment_id, indexers, *num_candidates)
+        .select(*deployment_id, indexers, *num_candidates, &context)
         .await
         .map_err(|err| JobError::Fatal(err.into()))?;
     if candidates.is_empty() {

bin/dipper-service/src/worker/handlers/selection_context.rs

@@ -0,0 +1,69 @@
+//! Shared utilities for gathering IISA selection context.
+
+use std::collections::HashMap;
+
+use dipper_iisa::{Indexer as IndexerCandidate, SelectionContext};
+use thegraph_core::{DeploymentId, IndexerId};
+
+use crate::{
+    registry::{AgreementRegistry, IndexingAgreementStatus},
+    worker::result::{JobError, JobResult},
+};
+
+/// Gather load balancing context for IISA selection.
+///
+/// This function queries the registry to build context about:
+/// - Which indexers already have active agreements for this deployment
+/// - What other deployments each candidate indexer is currently working on
+pub async fn gather_selection_context<R>(
+    registry: &R,
+    deployment_id: &DeploymentId,
+    candidates: &[IndexerCandidate],
+) -> JobResult<SelectionContext>
+where
+    R: AgreementRegistry,
+{
+    // Get indexers that already have active agreements for this deployment
+    let existing_indexers = registry
+        .get_indexing_agreements_by_deployment_id(deployment_id)
+        .await
+        .map_err(|err| JobError::Fatal(err.into()))?
+        .into_iter()
+        .filter(|a| is_active_agreement(&a.status))
+        .map(|a| a.indexer.id)
+        .collect::<Vec<_>>();
+
+    // Build pending agreements map for all candidates in a single batch query
+    // This tells IISA what other work each candidate is currently handling
+    let candidate_ids: Vec<IndexerId> = candidates.iter().map(|c| c.id).collect();
+
+    let all_agreements = registry
+        .get_active_indexing_agreements_by_indexer_ids(&candidate_ids)
+        .await
+        .map_err(|err| JobError::Fatal(err.into()))?;
+
+    // Group agreements by indexer ID
+    let mut pending_agreements: HashMap<IndexerId, Vec<DeploymentId>> = HashMap::new();
+    for agreement in all_agreements {
+        pending_agreements
+            .entry(agreement.indexer.id)
+            .or_default()
+            .push(agreement.voucher.metadata.subgraph_deployment_id);
+    }
+
+    Ok(SelectionContext {
+        existing_indexers,
+        pending_agreements,
+    })
+}
+
+/// Check if an agreement status represents an active agreement.
+///
+/// Active agreements are those that are either pending acceptance (Created)
+/// or currently in effect (Accepted).
+fn is_active_agreement(status: &IndexingAgreementStatus) -> bool {
+    matches!(
+        status,
+        IndexingAgreementStatus::Created | IndexingAgreementStatus::Accepted { .. }
+    )
+}

dipper-iisa/src/api.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use async_trait::async_trait;
 use thegraph_core::{DeploymentId, IndexerId};
 use url::Url;
@@ -11,6 +13,24 @@ pub struct Indexer {
     pub url: Url,
 }
 
+/// Context for load balancing during indexer selection.
+///
+/// This context provides the IISA with information about current system state,
+/// enabling intelligent load balancing decisions rather than naive selection.
+#[derive(Debug, Clone, Default)]
+pub struct SelectionContext {
+    /// Indexer IDs that already have active agreements for this deployment.
+    ///
+    /// Used to avoid selecting indexers that are already working on the same deployment.
+    pub existing_indexers: Vec<IndexerId>,
+
+    /// For each indexer, the deployments they have pending/active agreements for.
+    ///
+    /// Used to balance load across indexers by considering their current workload.
+    /// Key: Indexer ID, Value: List of deployment IDs they are working on.
+    pub pending_agreements: HashMap<IndexerId, Vec<DeploymentId>>,
+}
+
 /// The `SelectionError` enum represents the errors that can occur during the candidate selection
 /// process.
 #[derive(Debug, thiserror::Error)]
@@ -31,17 +51,30 @@ pub enum SelectionError {
 #[async_trait]
 pub trait CandidateSelection {
     /// Select one indexer from the given list of candidates.
+    ///
+    /// # Arguments
+    /// * `deployment_id` - The deployment to select an indexer for
+    /// * `candidates` - List of candidate indexers to choose from
+    /// * `context` - Load balancing context with existing assignments and pending work
     async fn select_one(
         &self,
         deployment_id: DeploymentId,
         candidates: Vec<Indexer>,
+        context: &SelectionContext,
     ) -> Result<Option<Indexer>, SelectionError>;
 
     /// Selects the best `num_candidates` indexers from the given list of candidates.
+    ///
+    /// # Arguments
+    /// * `deployment_id` - The deployment to select indexers for
+    /// * `candidates` - List of candidate indexers to choose from
+    /// * `num_candidates` - Maximum number of indexers to select
+    /// * `context` - Load balancing context with existing assignments and pending work
     async fn select(
         &self,
         deployment_id: DeploymentId,
         candidates: Vec<Indexer>,
         num_candidates: usize,
+        context: &SelectionContext,
     ) -> Result<Vec<Indexer>, SelectionError>;
 }

dipper-iisa/src/http_client.rs

@@ -16,7 +16,7 @@ use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use thegraph_core::{DeploymentId, IndexerId};
 
-use crate::api::{CandidateSelection, Indexer, SelectionError};
+use crate::api::{CandidateSelection, Indexer, SelectionContext, SelectionError};
 
 /// HTTP client for the IISA container service.
 #[derive(Clone)]
@@ -108,6 +108,47 @@ impl HttpIisaClient {
             url: indexer.url.to_string(),
         }
     }
+
+    /// Format existing indexers from context for the HTTP request.
+    ///
+    /// Returns `None` if the list is empty to skip serialization.
+    fn format_existing_indexers(context: &SelectionContext) -> Option<Vec<String>> {
+        if context.existing_indexers.is_empty() {
+            None
+        } else {
+            Some(
+                context
+                    .existing_indexers
+                    .iter()
+                    .map(|id| format!("{:#x}", id))
+                    .collect(),
+            )
+        }
+    }
+
+    /// Format pending agreements from context for the HTTP request.
+    ///
+    /// Returns `None` if the map is empty to skip serialization.
+    fn format_pending_agreements(
+        context: &SelectionContext,
+    ) -> Option<HashMap<String, Vec<String>>> {
+        if context.pending_agreements.is_empty() {
+            None
+        } else {
+            Some(
+                context
+                    .pending_agreements
+                    .iter()
+                    .map(|(indexer_id, deployment_ids)| {
+                        (
+                            format!("{:#x}", indexer_id),
+                            deployment_ids.iter().map(|d| d.to_string()).collect(),
+                        )
+                    })
+                    .collect(),
+            )
+        }
+    }
 }
 
 #[async_trait]
@@ -116,6 +157,7 @@ impl CandidateSelection for HttpIisaClient {
         &self,
         deployment_id: DeploymentId,
         candidates: Vec<Indexer>,
+        context: &SelectionContext,
     ) -> Result<Option<Indexer>, SelectionError> {
         if candidates.is_empty() {
             return Ok(None);
@@ -124,8 +166,8 @@ impl CandidateSelection for HttpIisaClient {
         let request = SelectionRequest {
             deployment_id: deployment_id.to_string(),
             candidates: Some(candidates.iter().map(Self::to_candidate).collect()),
-            existing_indexers: None,
-            pending_agreements: None,
+            existing_indexers: Self::format_existing_indexers(context),
+            pending_agreements: Self::format_pending_agreements(context),
             num_candidates: None,
         };
 
@@ -168,6 +210,7 @@ impl CandidateSelection for HttpIisaClient {
         deployment_id: DeploymentId,
         candidates: Vec<Indexer>,
         num_candidates: usize,
+        context: &SelectionContext,
     ) -> Result<Vec<Indexer>, SelectionError> {
         if candidates.is_empty() || num_candidates == 0 {
             return Ok(Vec::new());
@@ -176,8 +219,8 @@ impl CandidateSelection for HttpIisaClient {
         let request = SelectionRequest {
             deployment_id: deployment_id.to_string(),
             candidates: Some(candidates.iter().map(Self::to_candidate).collect()),
-            existing_indexers: None,
-            pending_agreements: None,
+            existing_indexers: Self::format_existing_indexers(context),
+            pending_agreements: Self::format_pending_agreements(context),
             num_candidates: Some(num_candidates),
         };
 
@@ -235,4 +278,64 @@ mod tests {
         let client = HttpIisaClient::new("http://localhost:8080/".to_string());
         assert_eq!(client.endpoint, "http://localhost:8080/");
     }
+
+    #[test]
+    fn test_format_existing_indexers_empty() {
+        let context = SelectionContext::default();
+        assert_eq!(HttpIisaClient::format_existing_indexers(&context), None);
+    }
+
+    #[test]
+    fn test_format_existing_indexers_with_data() {
+        let indexer_id: IndexerId = "0x1234567890123456789012345678901234567890"
+            .parse()
+            .unwrap();
+        let context = SelectionContext {
+            existing_indexers: vec![indexer_id],
+            pending_agreements: HashMap::new(),
+        };
+
+        let result = HttpIisaClient::format_existing_indexers(&context);
+        assert!(result.is_some());
+        let indexers = result.unwrap();
+        assert_eq!(indexers.len(), 1);
+        assert_eq!(indexers[0], "0x1234567890123456789012345678901234567890");
+    }
+
+    #[test]
+    fn test_format_pending_agreements_empty() {
+        let context = SelectionContext::default();
+        assert_eq!(HttpIisaClient::format_pending_agreements(&context), None);
+    }
+
+    #[test]
+    fn test_format_pending_agreements_with_data() {
+        let indexer_id: IndexerId = "0x1234567890123456789012345678901234567890"
+            .parse()
+            .unwrap();
+        let deployment_id: DeploymentId = "QmYwAPJzv5CZsnA625s3Xf2nemtYgPpHdWEz79ojWnPbdG"
+            .parse()
+            .unwrap();
+
+        let mut pending = HashMap::new();
+        pending.insert(indexer_id, vec![deployment_id]);
+
+        let context = SelectionContext {
+            existing_indexers: vec![],
+            pending_agreements: pending,
+        };
+
+        let result = HttpIisaClient::format_pending_agreements(&context);
+        assert!(result.is_some());
+        let agreements = result.unwrap();
+        assert_eq!(agreements.len(), 1);
+        assert!(agreements.contains_key("0x1234567890123456789012345678901234567890"));
+    }
+
+    #[test]
+    fn test_selection_context_default() {
+        let context = SelectionContext::default();
+        assert!(context.existing_indexers.is_empty());
+        assert!(context.pending_agreements.is_empty());
+    }
 }

dipper-iisa/src/lib.rs

@@ -1,5 +1,5 @@
 mod api;
 pub mod http_client;
 
-pub use api::{CandidateSelection, Indexer, SelectionError};
+pub use api::{CandidateSelection, Indexer, SelectionContext, SelectionError};
 pub use http_client::HttpIisaClient;