Skip to content

Commit 4f8922e

Browse files
jakubnojakubno
committed
Refactor pausing to sqlc
1 parent a3363e0 commit 4f8922e

File tree

4 files changed

+262
-149
lines changed

4 files changed

+262
-149
lines changed

‎packages/api/internal/orchestrator/pause_instance.go‎

Lines changed: 18 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,15 @@ import (
66
"fmt"
77
"time"
88

9+
"github.com/e2b-dev/infra/packages/shared/pkg/id"
910
"github.com/gogo/status"
11+
"github.com/jackc/pgx/v5/pgtype"
1012
"google.golang.org/grpc/codes"
1113

1214
"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
1315
"github.com/e2b-dev/infra/packages/api/internal/sandbox"
1416
"github.com/e2b-dev/infra/packages/db/queries"
1517
"github.com/e2b-dev/infra/packages/db/types"
16-
"github.com/e2b-dev/infra/packages/shared/pkg/db"
1718
"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
1819
"github.com/e2b-dev/infra/packages/shared/pkg/models/envbuild"
1920
"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
@@ -29,39 +30,36 @@ func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node,
2930
ctx, span := tracer.Start(ctx, "pause-sandbox")
3031
defer span.End()
3132

32-
snapshotConfig := &db.SnapshotInfo{
33+
result, err := o.sqlcDB.UpsertSnapshotEnvAndBuild(ctx, queries.UpsertSnapshotEnvAndBuildParams{
3334
BaseTemplateID: sbx.BaseTemplateID,
35+
TemplateID: id.Generate(),
36+
TeamID: sbx.TeamID,
3437
SandboxID: sbx.SandboxID,
35-
SandboxStartedAt: sbx.StartTime,
36-
VCPU: sbx.VCpu,
37-
RAMMB: sbx.RamMB,
38-
TotalDiskSizeMB: sbx.TotalDiskSizeMB,
38+
StartedAt: pgtype.Timestamptz{Time: sbx.StartTime, Valid: true},
39+
Status: string(envbuild.StatusSnapshotting),
40+
OriginNodeID: node.ID,
41+
Vcpu: sbx.VCpu,
42+
RamMb: sbx.RamMB,
43+
TotalDiskSizeMb: &sbx.TotalDiskSizeMB,
3944
Metadata: sbx.Metadata,
4045
KernelVersion: sbx.KernelVersion,
4146
FirecrackerVersion: sbx.FirecrackerVersion,
42-
EnvdVersion: sbx.EnvdVersion,
43-
EnvdSecured: sbx.EnvdAccessToken != nil,
47+
EnvdVersion: &sbx.EnvdVersion,
48+
Secure: sbx.EnvdAccessToken != nil,
4449
AllowInternetAccess: sbx.AllowInternetAccess,
4550
AutoPause: sbx.AutoPause,
4651
Config: &types.PausedSandboxConfig{
4752
Version: types.PausedSandboxConfigVersion,
4853
Network: sbx.Network,
4954
},
50-
}
51-
52-
envBuild, err := o.dbClient.NewSnapshotBuild(
53-
ctx,
54-
snapshotConfig,
55-
sbx.TeamID,
56-
sbx.NodeID,
57-
)
55+
})
5856
if err != nil {
59-
telemetry.ReportCriticalError(ctx, "error pausing sandbox", err)
57+
telemetry.ReportCriticalError(ctx, "error inserting snapshot for env", err)
6058

6159
return err
6260
}
6361

64-
err = snapshotInstance(ctx, o, node, sbx, envBuild.EnvID, envBuild.ID.String())
62+
err = snapshotInstance(ctx, o, node, sbx, result.TemplateID, result.BuildID.String())
6563
if errors.Is(err, PauseQueueExhaustedError{}) {
6664
telemetry.ReportCriticalError(ctx, "pause queue exhausted", err)
6765

@@ -79,8 +77,8 @@ func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node,
7977
Status: string(envbuild.StatusSuccess),
8078
FinishedAt: &now,
8179
Reason: types.BuildReason{},
82-
BuildID: envBuild.ID,
83-
TemplateID: envBuild.EnvID,
80+
BuildID: result.BuildID,
81+
TemplateID: result.TemplateID,
8482
})
8583
if err != nil {
8684
telemetry.ReportCriticalError(ctx, "error pausing sandbox", err)

‎packages/db/queries/create_new_snapshot.sql.go‎

Lines changed: 156 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/db/queries/snapshots/create_new_snapshot.sql‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
-- name: UpsertSnapshotEnvAndBuild :one
2+
WITH updated_snapshot AS (
3+
UPDATE "public"."snapshots" s
4+
SET
5+
metadata = @metadata,
6+
sandbox_started_at = @started_at,
7+
origin_node_id = @origin_node_id,
8+
auto_pause = @auto_pause,
9+
config = @config
10+
FROM "public"."envs" e
11+
WHERE
12+
s.sandbox_id = @sandbox_id
13+
AND e.id = s.env_id
14+
AND e.team_id = @team_id
15+
RETURNING s.env_id
16+
),
17+
inserted_env AS (
18+
INSERT INTO "public"."envs" (id, public, created_by, team_id, updated_at)
19+
SELECT @template_id, FALSE, NULL, @team_id, now()
20+
WHERE NOT EXISTS (SELECT 1 FROM updated_snapshot)
21+
ON CONFLICT (id) DO NOTHING
22+
RETURNING id AS env_id
23+
),
24+
inserted_snapshot AS (
25+
INSERT INTO "public"."snapshots" (
26+
sandbox_id,
27+
base_env_id,
28+
team_id,
29+
env_id,
30+
metadata,
31+
sandbox_started_at,
32+
env_secure,
33+
allow_internet_access,
34+
origin_node_id,
35+
auto_pause,
36+
config
37+
)
38+
SELECT
39+
@sandbox_id,
40+
@base_template_id,
41+
@team_id,
42+
COALESCE(
43+
(SELECT env_id FROM inserted_env LIMIT 1),
44+
@template_id -- env already existed
45+
) AS env_id,
46+
@metadata,
47+
@started_at,
48+
@secure,
49+
@allow_internet_access,
50+
@origin_node_id,
51+
@auto_pause,
52+
@config
53+
WHERE NOT EXISTS (SELECT 1 FROM updated_snapshot)
54+
RETURNING env_id
55+
),
56+
final_env AS (
57+
-- If we updated an existing snapshot, use that env_id.
58+
-- Otherwise use env_id from the newly inserted snapshot.
59+
SELECT env_id FROM updated_snapshot
60+
UNION ALL
61+
SELECT env_id FROM inserted_snapshot
62+
)
63+
INSERT INTO "public"."env_builds" (
64+
env_id,
65+
vcpu,
66+
ram_mb,
67+
free_disk_size_mb,
68+
kernel_version,
69+
firecracker_version,
70+
envd_version,
71+
status,
72+
cluster_node_id,
73+
total_disk_size_mb,
74+
updated_at
75+
)
76+
SELECT
77+
(SELECT env_id FROM final_env LIMIT 1),
78+
@vcpu,
79+
@ram_mb,
80+
0,
81+
@kernel_version,
82+
@firecracker_version,
83+
@envd_version,
84+
@status,
85+
@origin_node_id,
86+
@total_disk_size_mb,
87+
now()
88+
RETURNING id as build_id, env_id as template_id;

0 commit comments

Comments
 (0)