feat: workspace keepalive task

Author: alemidev

src/client.rs

@@ -171,26 +171,49 @@ impl Client {
 	/// Join and return a [`Workspace`].
 	#[tracing::instrument(skip(self, workspace), fields(ws = %workspace))]
 	pub async fn attach_workspace(&self, workspace: uuid::Uuid) -> ConnectionResult<Workspace> {
-		let token = self
-			.0
-			.session
-			.clone()
-			.access_workspace(WorkspaceRequest {
+		let mut session_client = self.0.session.clone();
+		let token = session_client
+			.get_workspace_token(WorkspaceRequest {
 				id: Identifier::from(workspace),
 			})
 			.await?
 			.into_inner();
 
+		let workspace_claims = InternallyMutable::new(token);
+
 		let ws = Workspace::connect(
 			workspace,
 			self.0.user.clone(),
 			self.0.config.clone(),
-			token,
+			workspace_claims.channel(),
 			self.0.claims.channel(),
 		)
 		.await?;
-
 		self.0.workspaces.insert(workspace, ws.clone());
+		let mut workspace_client = ws.services().ws();
+
+		let weak = Arc::downgrade(&ws.0);
+		tokio::spawn(async move {
+			let fut = async move {
+				loop {
+					// TODO either configurable token refresh time or calculate depending on token lifetime
+					tokio::time::sleep(std::time::Duration::from_secs(240)).await;
+					if weak.upgrade().is_none() { break };
+					let new_credentials = session_client.get_workspace_token(
+						tonic::Request::new(WorkspaceRequest { id: Identifier::from(workspace) })
+					)
+						.await?
+						.into_inner();
+					workspace_claims.set(new_credentials);
+					workspace_client.keep_alive(tonic::Request::new(Empty {})).await?;
+				}
+				Ok::<(), tonic::Status>(())
+			};
+
+			if let Err(e) = fut.await {
+				tracing::error!("error in keepalive task for workspace {workspace}: {e}");
+			}
+		});
 
 		Ok(ws)
 	}

src/network.rs

@@ -40,10 +40,10 @@ impl Services {
 		let channel = Endpoint::from_shared(dest.to_string())?.connect().await?;
 		let inter = WorkspaceInterceptor { session, workspace };
 		Ok(Self {
-			cursor: CursorClient::with_interceptor(channel.clone(), inter.clone()),
 			workspace: WorkspaceClient::with_interceptor(channel.clone(), inter.clone()),
-			// TODO technically we could keep buffers on separate servers, and thus manage buffer
-			// connections separately, but for now it's more convenient to bundle them with workspace
+			// TODO technically we could keep buffers and cursors on separate servers, and thus manage
+			// their connections separately, but for now it's more convenient to bundle them with workspace
+			cursor: CursorClient::with_interceptor(channel.clone(), inter.clone()),
 			buffer: BufferClient::with_interceptor(channel.clone(), inter.clone()),
 		})
 	}

src/workspace.rs

@@ -45,7 +45,7 @@ use napi_derive::napi;
 #[derive(Debug, Clone)]
 #[cfg_attr(feature = "py", pyo3::pyclass)]
 #[cfg_attr(feature = "js", napi)]
-pub struct Workspace(Arc<WorkspaceInner>);
+pub struct Workspace(pub(crate) Arc<WorkspaceInner>);
 
 #[derive(Debug)]
 struct WorkspaceInner {
@@ -87,17 +87,16 @@ impl AsyncReceiver<Event> for Workspace {
 }
 
 impl Workspace {
-	#[tracing::instrument(skip(id, user, token, claims), fields(ws = %id))]
+	#[tracing::instrument(skip(id, user, workspace_claim, user_claim), fields(ws = %id))]
 	pub(crate) async fn connect(
 		id: Uuid,
 		user: Arc<User>,
 		config: crate::api::Config,
-		token: Token,
-		claims: tokio::sync::watch::Receiver<codemp_proto::common::Token>,
+		workspace_claim: tokio::sync::watch::Receiver<codemp_proto::common::Token>,
+		user_claim: tokio::sync::watch::Receiver<codemp_proto::common::Token>,
 	) -> ConnectionResult<Self> {
-		let workspace_claim = InternallyMutable::new(token);
 		let services =
-			Services::try_new(&config.endpoint(), claims, workspace_claim.channel()).await?;
+			Services::try_new(&config.endpoint(), user_claim, workspace_claim).await?;
 		let ws_stream = services.ws().attach(Empty {}).await?.into_inner();
 
 		let (tx, rx) = mpsc::channel(128);
@@ -146,6 +145,10 @@ impl Workspace {
 		Ok(ws)
 	}
 
+	pub(crate) fn services(&self) -> &Services {
+		&self.0.services
+	}
+
 	/// drop arc, return true if was last
 	pub(crate) fn consume(self) -> bool {
 		Arc::into_inner(self.0).is_some()