mz_storage/
server.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! An interactive dataflow server.
11
12use std::sync::{Arc, Mutex};
13use std::task::{Context, Poll};
14
15use mz_cluster::client::{ClusterClient, ClusterSpec};
16use mz_cluster_client::client::TimelyConfig;
17use mz_ore::metrics::MetricsRegistry;
18use mz_ore::now::NowFn;
19use mz_ore::tracing::TracingHandle;
20use mz_persist_client::cache::PersistClientCache;
21use mz_rocksdb::config::SharedWriteBufferManager;
22use mz_service::client::GenericClient;
23use mz_storage_client::client::{StorageCommand, StorageResponse};
24use mz_storage_types::connections::ConnectionContext;
25use mz_txn_wal::operator::TxnsContext;
26use timely::worker::Worker as TimelyWorker;
27use tokio::sync::mpsc;
28use uuid::Uuid;
29
30use crate::metrics::StorageMetrics;
31use crate::storage_state::{StorageInstanceContext, Worker};
32
33/// Configures a dataflow server.
34#[derive(Clone)]
35struct Config {
36    /// `persist` client cache.
37    pub persist_clients: Arc<PersistClientCache>,
38    /// Context necessary for rendering txn-wal operators.
39    pub txns_ctx: TxnsContext,
40    /// A process-global handle to tracing configuration.
41    pub tracing_handle: Arc<TracingHandle>,
42    /// Function to get wall time now.
43    pub now: NowFn,
44    /// Configuration for source and sink connection.
45    pub connection_context: ConnectionContext,
46    /// Other configuration for storage instances.
47    pub instance_context: StorageInstanceContext,
48
49    /// Metrics for storage
50    pub metrics: StorageMetrics,
51    /// Shared rocksdb write buffer manager
52    pub shared_rocksdb_write_buffer_manager: SharedWriteBufferManager,
53}
54
55/// Initiates a timely dataflow computation, processing storage commands.
56pub async fn serve(
57    timely_config: TimelyConfig,
58    metrics_registry: &MetricsRegistry,
59    persist_clients: Arc<PersistClientCache>,
60    txns_ctx: TxnsContext,
61    tracing_handle: Arc<TracingHandle>,
62    now: NowFn,
63    connection_context: ConnectionContext,
64    instance_context: StorageInstanceContext,
65) -> Result<impl Fn() -> StorageServerClient + use<>, anyhow::Error> {
66    let config = Config {
67        persist_clients,
68        txns_ctx,
69        tracing_handle,
70        now,
71        connection_context,
72        instance_context,
73        metrics: StorageMetrics::register_with(metrics_registry),
74        // The shared RocksDB `WriteBufferManager` is shared between the workers.
75        // It protects (behind a shared mutex) a `Weak` that will be upgraded and shared when the
76        // first worker attempts to initialize it.
77        shared_rocksdb_write_buffer_manager: Default::default(),
78    };
79    let tokio_executor = tokio::runtime::Handle::current();
80
81    let timely_container = config.build_cluster(timely_config, tokio_executor).await?;
82    let timely_container = Arc::new(Mutex::new(timely_container));
83
84    let client_builder =
85        move || StorageServerClient(ClusterClient::new(Arc::clone(&timely_container)));
86
87    Ok(client_builder)
88}
89
90/// A client to the Timely cluster of a storage server.
91#[derive(Debug)]
92pub struct StorageServerClient(ClusterClient<Config>);
93
94impl GenericClient<StorageCommand, StorageResponse> for StorageServerClient {
95    fn send(
96        &mut self,
97        cmd: StorageCommand,
98    ) -> impl Future<Output = Result<(), anyhow::Error>> + Send {
99        self.0.send(cmd)
100    }
101
102    fn poll_recv(
103        &mut self,
104        cx: &mut Context<'_>,
105    ) -> Poll<Result<Option<StorageResponse>, anyhow::Error>> {
106        self.0.poll_recv(cx)
107    }
108}
109
110impl ClusterSpec for Config {
111    type Command = StorageCommand;
112    type Response = StorageResponse;
113
114    const NAME: &str = "storage";
115
116    fn run_worker(
117        &self,
118        timely_worker: &mut TimelyWorker,
119        client_rx: mpsc::UnboundedReceiver<(
120            Uuid,
121            mpsc::UnboundedReceiver<StorageCommand>,
122            mpsc::UnboundedSender<StorageResponse>,
123        )>,
124    ) {
125        Worker::new(
126            timely_worker,
127            client_rx,
128            self.metrics.clone(),
129            self.now.clone(),
130            self.connection_context.clone(),
131            self.instance_context.clone(),
132            Arc::clone(&self.persist_clients),
133            self.txns_ctx.clone(),
134            Arc::clone(&self.tracing_handle),
135            self.shared_rocksdb_write_buffer_manager.clone(),
136        )
137        .run();
138    }
139}
mz_storage/server.rs

mz_storage/
server.rs
