Merge pull request #998 from geo-engine/quota_log_wip

Quota and Data usage Logging

Author: michaelmattig

datatypes/src/primitives/datetime.rs

@@ -322,6 +322,12 @@ impl From<chrono::DateTime<chrono::Utc>> for DateTime {
     }
 }
 
+impl From<DateTime> for chrono::DateTime<chrono::Utc> {
+    fn from(datetime: DateTime) -> Self {
+        datetime.datetime
+    }
+}
+
 impl From<DateTime> for chrono::DateTime<chrono::FixedOffset> {
     fn from(datetime: DateTime) -> Self {
         Self::from(&datetime)
@@ -334,12 +340,6 @@ impl From<&DateTime> for chrono::DateTime<chrono::FixedOffset> {
     }
 }
 
-impl From<DateTime> for chrono::DateTime<chrono::Utc> {
-    fn from(datetime: DateTime) -> Self {
-        Self::from(&datetime)
-    }
-}
-
 impl From<&DateTime> for chrono::DateTime<chrono::Utc> {
     fn from(datetime: &DateTime) -> Self {
         datetime.datetime

operators/benches/query_chunks.rs

@@ -25,7 +25,7 @@ use geoengine_datatypes::{
         SpatialPartition2D, SpatialResolution, TimeInterval, VectorQueryRectangle,
     },
     raster::Pixel,
-    util::{test::TestDefault, Identifier},
+    util::test::TestDefault,
 };
 use geoengine_operators::{
     engine::{
@@ -34,7 +34,7 @@ use geoengine_operators::{
         SingleVectorMultipleRasterSources, StatisticsWrappingMockExecutionContext,
         TypedRasterQueryProcessor, VectorOperator, VectorQueryProcessor, WorkflowOperatorPath,
     },
-    meta::quota::{ComputationContext, ComputationUnit, QuotaCheck, QuotaChecker, QuotaTracking},
+    meta::quota::{QuotaCheck, QuotaChecker, QuotaTracking},
     processing::{
         AggregateFunctionParams, ColumnNames, FeatureAggregationMethod, NeighborhoodAggregate,
         NeighborhoodAggregateParams, NeighborhoodParams, RasterVectorJoin, RasterVectorJoinParams,
@@ -70,17 +70,18 @@ type BenchmarkElementCounts = HashMap<String, u64>;
 fn setup_contexts() -> (StatisticsWrappingMockExecutionContext, MockQueryContext) {
     let exe_ctx = StatisticsWrappingMockExecutionContext::test_default();
 
-    let computation_unit = ComputationUnit {
-        issuer: uuid::Uuid::new_v4(),
-        context: ComputationContext::new(),
-    };
+    let user = uuid::Uuid::new_v4();
+    let workflow = uuid::Uuid::new_v4();
+    let computation = uuid::Uuid::new_v4();
 
     let query_ctx = MockQueryContext::new_with_query_extensions(
         ChunkByteSize::test_default(),
         None,
         Some(QuotaTracking::new(
             tokio::sync::mpsc::unbounded_channel().0,
-            computation_unit,
+            user,
+            workflow,
+            computation,
         )),
         Some(Box::new(MockQuotaChecker) as QuotaChecker),
     );

operators/src/adapters/stream_statistics_adapter.rs

@@ -16,6 +16,9 @@ pub struct StreamStatisticsAdapter<S> {
     span: Span,
     quota: QuotaTracking,
     path: WorkflowOperatorPath,
+    operator_name: &'static str,
+    /// If the wrapped stream is a source, pass the value of the `data` parameter to the quota tracking
+    data: Option<String>,
 }
 
 impl<S> StreamStatisticsAdapter<S> {
@@ -24,6 +27,8 @@ impl<S> StreamStatisticsAdapter<S> {
         span: Span,
         quota: QuotaTracking,
         path: WorkflowOperatorPath,
+        operator_name: &'static str,
+        data: Option<String>,
     ) -> StreamStatisticsAdapter<S> {
         StreamStatisticsAdapter {
             stream,
@@ -32,6 +37,8 @@ impl<S> StreamStatisticsAdapter<S> {
             span,
             quota,
             path,
+            operator_name,
+            data,
         }
     }
 
@@ -76,7 +83,11 @@ where
                     empty = false,
                 );
 
-                (*this.quota).work_unit_done();
+                (*this.quota).work_unit_done(
+                    this.operator_name,
+                    this.path.clone(),
+                    this.data.clone(),
+                );
             }
             None => {
                 tracing::debug!(
@@ -100,9 +111,8 @@ mod tests {
 
     use super::*;
 
-    use crate::meta::quota::{ComputationContext, ComputationUnit, QuotaMessage};
+    use crate::meta::quota::{ComputationUnit, QuotaMessage};
     use futures::StreamExt;
-    use geoengine_datatypes::util::Identifier;
     use tokio::sync::mpsc::unbounded_channel;
     use tracing::{span, Level};
     use uuid::Uuid;
@@ -112,14 +122,17 @@ mod tests {
         let v = vec![1, 2, 3];
         let v_stream = futures::stream::iter(v);
         let (tx, mut rx) = unbounded_channel::<QuotaMessage>();
-        let issuer = Uuid::new_v4();
-        let context = ComputationContext::new();
-        let quota = QuotaTracking::new(tx, ComputationUnit { issuer, context });
+        let user = Uuid::new_v4();
+        let workflow = Uuid::new_v4();
+        let computation = Uuid::new_v4();
+        let quota = QuotaTracking::new(tx, user, workflow, computation);
         let mut v_stat_stream = StreamStatisticsAdapter::new(
             v_stream,
             span!(Level::TRACE, "test"),
             quota,
             WorkflowOperatorPath::initialize_root(),
+            "test",
+            None,
         );
 
         let one = v_stat_stream.next().await;
@@ -142,7 +155,15 @@ mod tests {
 
         assert_eq!(
             rx.recv().await.unwrap(),
-            ComputationUnit { issuer, context }.into()
+            ComputationUnit {
+                user,
+                workflow,
+                computation,
+                operator_name: "test",
+                operator_path: WorkflowOperatorPath::initialize_root(),
+                data: None,
+            }
+            .into()
         );
     }
 }

operators/src/cache/cache_operator.rs

@@ -6,7 +6,7 @@ use crate::cache::shared_cache::{AsyncCache, SharedCache};
 use crate::engine::{
     CanonicOperatorName, ChunkByteSize, InitializedRasterOperator, InitializedVectorOperator,
     QueryContext, QueryProcessor, RasterResultDescriptor, ResultDescriptor,
-    TypedRasterQueryProcessor,
+    TypedRasterQueryProcessor, WorkflowOperatorPath,
 };
 use crate::error::Error;
 use crate::util::Result;
@@ -90,6 +90,14 @@ impl InitializedRasterOperator for InitializedCacheOperator<Box<dyn InitializedR
     fn canonic_name(&self) -> CanonicOperatorName {
         self.source.canonic_name()
     }
+
+    fn name(&self) -> &'static str {
+        self.source.name()
+    }
+
+    fn path(&self) -> WorkflowOperatorPath {
+        self.source.path()
+    }
 }
 
 impl InitializedVectorOperator for InitializedCacheOperator<Box<dyn InitializedVectorOperator>> {
@@ -137,6 +145,14 @@ impl InitializedVectorOperator for InitializedCacheOperator<Box<dyn InitializedV
     fn canonic_name(&self) -> CanonicOperatorName {
         self.source.canonic_name()
     }
+
+    fn name(&self) -> &'static str {
+        self.source.name()
+    }
+
+    fn path(&self) -> WorkflowOperatorPath {
+        self.source.path()
+    }
 }
 
 /// A cache operator that caches the results of its source operator

operators/src/engine/execution_context.rs

@@ -1,7 +1,7 @@
 use super::query::QueryAbortRegistration;
 use super::{
     CreateSpan, InitializedPlotOperator, InitializedRasterOperator, InitializedVectorOperator,
-    MockQueryContext, WorkflowOperatorPath,
+    MockQueryContext,
 };
 use crate::engine::{
     ChunkByteSize, RasterResultDescriptor, ResultDescriptor, VectorResultDescriptor,
@@ -40,21 +40,18 @@ pub trait ExecutionContext: Send
         &self,
         op: Box<dyn InitializedRasterOperator>,
         span: CreateSpan,
-        path: WorkflowOperatorPath, // TODO: remove and allow operators to tell its path
     ) -> Box<dyn InitializedRasterOperator>;
 
     fn wrap_initialized_vector_operator(
         &self,
         op: Box<dyn InitializedVectorOperator>,
         span: CreateSpan,
-        path: WorkflowOperatorPath,
     ) -> Box<dyn InitializedVectorOperator>;
 
     fn wrap_initialized_plot_operator(
         &self,
         op: Box<dyn InitializedPlotOperator>,
         span: CreateSpan,
-        path: WorkflowOperatorPath,
     ) -> Box<dyn InitializedPlotOperator>;
 
     async fn resolve_named_data(&self, data: &NamedData) -> Result<DataId>;
@@ -187,7 +184,6 @@ impl ExecutionContext for MockExecutionContext {
         &self,
         op: Box<dyn InitializedRasterOperator>,
         _span: CreateSpan,
-        _path: WorkflowOperatorPath,
     ) -> Box<dyn InitializedRasterOperator> {
         op
     }
@@ -196,7 +192,6 @@ impl ExecutionContext for MockExecutionContext {
         &self,
         op: Box<dyn InitializedVectorOperator>,
         _span: CreateSpan,
-        _path: WorkflowOperatorPath,
     ) -> Box<dyn InitializedVectorOperator> {
         op
     }
@@ -205,7 +200,6 @@ impl ExecutionContext for MockExecutionContext {
         &self,
         op: Box<dyn InitializedPlotOperator>,
         _span: CreateSpan,
-        _path: WorkflowOperatorPath,
     ) -> Box<dyn InitializedPlotOperator> {
         op
     }
@@ -382,25 +376,22 @@ impl ExecutionContext for StatisticsWrappingMockExecutionContext {
         &self,
         op: Box<dyn InitializedRasterOperator>,
         span: CreateSpan,
-        path: WorkflowOperatorPath,
     ) -> Box<dyn InitializedRasterOperator> {
-        InitializedOperatorWrapper::new(op, span, path).boxed()
+        InitializedOperatorWrapper::new(op, span).boxed()
     }
 
     fn wrap_initialized_vector_operator(
         &self,
         op: Box<dyn InitializedVectorOperator>,
         span: CreateSpan,
-        path: WorkflowOperatorPath,
     ) -> Box<dyn InitializedVectorOperator> {
-        InitializedOperatorWrapper::new(op, span, path).boxed()
+        InitializedOperatorWrapper::new(op, span).boxed()
     }
 
     fn wrap_initialized_plot_operator(
         &self,
         op: Box<dyn InitializedPlotOperator>,
         _span: CreateSpan,
-        _path: WorkflowOperatorPath,
     ) -> Box<dyn InitializedPlotOperator> {
         op
     }

operators/src/engine/mod.rs

@@ -195,7 +195,7 @@ macro_rules! span_fn {
                 tracing::Level::TRACE,
                 <$op>::TYPE_NAME,
                 path = %path,
-                query_counter = %query_counter
+                query_counter = %query_counter,
             )
         }
     };

operators/src/engine/operator.rs

@@ -51,7 +51,7 @@ pub trait RasterOperator:
         #[allow(clippy::used_underscore_items)] // TODO: maybe rename?
         let op = self._initialize(path.clone(), context).await?;
 
-        Ok(context.wrap_initialized_raster_operator(op, span, path))
+        Ok(context.wrap_initialized_raster_operator(op, span))
     }
 
     /// Wrap a box around a `RasterOperator`
@@ -86,7 +86,7 @@ pub trait VectorOperator:
         let span = self.span();
         debug!("Initialize {}, path: {}", self.typetag_name(), &path);
         let op = self._initialize(path.clone(), context).await?;
-        Ok(context.wrap_initialized_vector_operator(op, span, path))
+        Ok(context.wrap_initialized_vector_operator(op, span))
     }
 
     /// Wrap a box around a `VectorOperator`
@@ -121,7 +121,7 @@ pub trait PlotOperator:
         debug!("Initialize {}, path: {}", self.typetag_name(), &path);
         #[allow(clippy::used_underscore_items)] // TODO: maybe rename?
         let op = self._initialize(path.clone(), context).await?;
-        Ok(context.wrap_initialized_plot_operator(op, span, path))
+        Ok(context.wrap_initialized_plot_operator(op, span))
     }
 
     /// Wrap a box around a `PlotOperator`
@@ -135,6 +135,7 @@ pub trait PlotOperator:
     fn span(&self) -> CreateSpan;
 }
 
+// TODO: implement a derive macro for common fields of operators: name, path, data, result_descriptor and automatically implement common trait functions
 pub trait InitializedRasterOperator: Send + Sync {
     /// Get the result descriptor of the `Operator`
     fn result_descriptor(&self) -> &RasterResultDescriptor;
@@ -152,6 +153,17 @@ pub trait InitializedRasterOperator: Send + Sync {
 
     /// Get a canonic representation of the operator and its sources
     fn canonic_name(&self) -> CanonicOperatorName;
+
+    /// Get the unique name of the operator
+    fn name(&self) -> &'static str;
+
+    // Get the path of the operator in the workflow
+    fn path(&self) -> WorkflowOperatorPath;
+
+    /// Return the name of the data loaded by the operator (if any)
+    fn data(&self) -> Option<String> {
+        None
+    }
 }
 
 pub trait InitializedVectorOperator: Send + Sync {
@@ -172,6 +184,17 @@ pub trait InitializedVectorOperator: Send + Sync {
     /// Get a canonic representation of the operator and its sources.
     /// This only includes *logical* operators, not wrappers
     fn canonic_name(&self) -> CanonicOperatorName;
+
+    /// Get the unique name of the operator
+    fn name(&self) -> &'static str;
+
+    // Get the path of the operator in the workflow
+    fn path(&self) -> WorkflowOperatorPath;
+
+    /// Return the name of the data loaded by the operator (if any)
+    fn data(&self) -> Option<String> {
+        None
+    }
 }
 
 /// A canonic name for an operator and its sources
@@ -247,6 +270,18 @@ impl InitializedRasterOperator for Box<dyn InitializedRasterOperator> {
     fn canonic_name(&self) -> CanonicOperatorName {
         self.as_ref().canonic_name()
     }
+
+    fn name(&self) -> &'static str {
+        self.as_ref().name()
+    }
+
+    fn path(&self) -> WorkflowOperatorPath {
+        self.as_ref().path()
+    }
+
+    fn data(&self) -> Option<String> {
+        self.as_ref().data()
+    }
 }
 
 impl InitializedVectorOperator for Box<dyn InitializedVectorOperator> {
@@ -261,6 +296,18 @@ impl InitializedVectorOperator for Box<dyn InitializedVectorOperator> {
     fn canonic_name(&self) -> CanonicOperatorName {
         self.as_ref().canonic_name()
     }
+
+    fn name(&self) -> &'static str {
+        self.as_ref().name()
+    }
+
+    fn path(&self) -> WorkflowOperatorPath {
+        self.as_ref().path()
+    }
+
+    fn data(&self) -> Option<String> {
+        self.as_ref().data()
+    }
 }
 
 impl InitializedPlotOperator for Box<dyn InitializedPlotOperator> {