[server] WorkspaceService.sendHeartBeat

Author: geropl

components/server/src/workspace/gitpod-server-impl.ts

@@ -96,7 +96,6 @@ import { WorkspaceManagerClientProvider } from "@gitpod/ws-manager/lib/client-pr
 import {
     AdmissionLevel,
     ControlAdmissionRequest,
-    MarkActiveRequest,
     StopWorkspacePolicy,
     TakeSnapshotRequest,
 } from "@gitpod/ws-manager/lib/core_pb";
@@ -173,7 +172,7 @@ import { RedisSubscriber } from "../messaging/redis-subscriber";
 import { UsageService } from "../orgs/usage-service";
 import { UserService } from "../user/user-service";
 import { SSHKeyService } from "../user/sshkey-service";
-import { StartWorkspaceOptions, WorkspaceService, mapGrpcError } from "./workspace-service";
+import { StartWorkspaceOptions, WorkspaceService } from "./workspace-service";
 import { GitpodTokenService } from "../user/gitpod-token-service";
 import { EnvVarService } from "../user/env-var-service";
 
@@ -1098,36 +1097,9 @@ export class GitpodServerImpl implements GitpodServerWithTracing, Disposable {
 
         const user = await this.checkAndBlockUser("sendHeartBeat", undefined, { instanceId });
 
-        try {
-            const wsi = await this.workspaceDb.trace(ctx).findInstanceById(instanceId);
-            if (!wsi) {
-                throw new ApplicationError(ErrorCodes.NOT_FOUND, "workspace does not exist");
-            }
-
-            const ws = await this.workspaceDb.trace(ctx).findById(wsi.workspaceId);
-            if (!ws) {
-                throw new ApplicationError(ErrorCodes.NOT_FOUND, "workspace does not exist");
-            }
-            await this.guardAccess({ kind: "workspaceInstance", subject: wsi, workspace: ws }, "update");
-
-            const wasClosed = !!(options && options.wasClosed);
-            await this.workspaceDb.trace(ctx).updateLastHeartbeat(instanceId, user.id, new Date(), wasClosed);
-
-            const req = new MarkActiveRequest();
-            req.setId(instanceId);
-            req.setClosed(wasClosed);
-
-            const client = await this.workspaceManagerClientProvider.get(wsi.region);
-            await client.markActive(ctx, req);
-        } catch (e) {
-            if (e.message && typeof e.message === "string" && (e.message as String).endsWith("does not exist")) {
-                // This is an old tab with open workspace: drop silently
-                return;
-            } else {
-                e = mapGrpcError(e);
-                throw e;
-            }
-        }
+        await this.workspaceService.sendHeartBeat(user.id, options, (instance, workspace) =>
+            this.guardAccess({ kind: "workspaceInstance", subject: instance, workspace }, "update"),
+        );
     }
 
     async getWorkspaceOwner(ctx: TraceContext, workspaceId: string): Promise<UserInfo | undefined> {

components/server/src/workspace/workspace-service.spec.db.ts

@@ -365,6 +365,19 @@ describe("WorkspaceService", async () => {
             ) => {},
         }); // returns without error in case of non-running workspace
     });
+
+    it("should sendHeartBeat", async () => {
+        const svc = container.get(WorkspaceService);
+        await createTestWorkspace(svc, org, owner, project);
+
+        await expectError(
+            ErrorCodes.NOT_FOUND,
+            svc.sendHeartBeat(owner.id, {
+                instanceId: "non-existing-instanceId",
+            }),
+            "should fail on non-running workspace",
+        );
+    });
 });
 
 async function createTestWorkspace(svc: WorkspaceService, org: Organization, owner: User, project: Project) {

components/server/src/workspace/workspace-service.ts

@@ -40,6 +40,7 @@ import {
     PortSpec,
     ControlPortRequest,
     SetTimeoutRequest,
+    MarkActiveRequest,
 } from "@gitpod/ws-manager/lib";
 import { WorkspaceStarter } from "./workspace-starter";
 import { LogContext, log } from "@gitpod/gitpod-protocol/lib/util/logging";
@@ -775,6 +776,43 @@ export class WorkspaceService {
             aborted.resolve(false);
         }
     }
+
+    public async sendHeartBeat(
+        userId: string,
+        options: GitpodServer.SendHeartBeatOptions,
+        check: (instance: WorkspaceInstance, workspace: Workspace) => Promise<void> = async () => {},
+    ): Promise<void> {
+        const instanceId = options.instanceId;
+        const instance = await this.db.findInstanceById(instanceId);
+        if (!instance) {
+            throw new ApplicationError(ErrorCodes.NOT_FOUND, "workspace does not exist");
+        }
+        const workspaceId = instance.workspaceId;
+        await this.auth.checkPermissionOnWorkspace(userId, "access", workspaceId);
+
+        try {
+            const workspace = await this.getWorkspace(userId, workspaceId);
+            await check(instance, workspace);
+
+            const wasClosed = !!(options && options.wasClosed);
+            await this.db.updateLastHeartbeat(instanceId, userId, new Date(), wasClosed);
+
+            const req = new MarkActiveRequest();
+            req.setId(instanceId);
+            req.setClosed(wasClosed);
+
+            const client = await this.clientProvider.get(instance.region);
+            await client.markActive({}, req);
+        } catch (e) {
+            if (e.message && typeof e.message === "string" && (e.message as String).endsWith("does not exist")) {
+                // This is an old tab with open workspace: drop silently
+                return;
+            } else {
+                e = mapGrpcError(e);
+                throw e;
+            }
+        }
+    }
 }
 
 // TODO(gpl) Make private after FGA rollout