Persist sandbox information locally in orchestrator

Makefile

@@ -31,7 +31,7 @@ tf_vars := TF_VAR_client_machine_type=$(CLIENT_MACHINE_TYPE) \
 	TF_VAR_template_bucket_location=$(TEMPLATE_BUCKET_LOCATION) \
 	TF_VAR_clickhouse_connection_string=$(CLICKHOUSE_CONNECTION_STRING) \
 	TF_VAR_clickhouse_username=$(CLICKHOUSE_USERNAME) \
-	TF_VAR_clickhouse_database=$(CLICKHOUSE_DATABASE) 
+	TF_VAR_clickhouse_database=$(CLICKHOUSE_DATABASE)
 
 # Login for Packer and Docker (uses gcloud user creds)
 # Login for Terraform (uses application default creds)
@@ -153,17 +153,24 @@ copy-public-builds:
 	gsutil cp -r gs://e2b-prod-public-builds/firecrackers/* gs://$(GCP_PROJECT_ID)-fc-versions/
 
 
+.PHONY: migrate
+migrate:
+	$(MAKE) -C packages/db migrate-postgres/up
+	# $(MAKE) -C packages/shared migrate-clickhouse/up
+
+
 .PHONY: generate
-generate: generate/api generate/orchestrator generate/template-manager generate/envd generate/db
+generate:generate/api
+generate:generate/orchestrator
+generate:generate/template-manager
+generate:generate/envd
+generate:generate/db
+generate:generate/orchestrator/internal
+generate:generate/shared
 generate/%:
-	@echo "Generating code for *$(notdir $@)*"
 	$(MAKE) -C packages/$(notdir $@) generate
 	@printf "\n\n"
 
-.PHONY: migrate
-migrate:
-	$(MAKE) -C packages/db migrate-postgres/up
-	# $(MAKE) -C packages/shared migrate-clickhouse/up
 
 .PHONY: switch-env
 switch-env:
@@ -221,4 +228,4 @@ grafana-apply:
 
 .PHONY: connect-orchestrator
 connect-orchestrator:
-	$(MAKE) -C tests/integration connect-orchestrator
+	$(MAKE) -C tests/integration connect-orchestrator

packages/orchestrator/Makefile

@@ -13,7 +13,7 @@ init:
 generate:
 	# You need to install protobuf (brew install protobuf) and following go packages: protoc-gen-go, protoc-gen-go-grpc
 	# https://grpc.io/docs/languages/go/quickstart/
-	@echo "Generating..."
+	@echo "Generating protoc..."
 	@protoc --go_out=../shared/pkg/grpc/orchestrator/ --go_opt=paths=source_relative --go-grpc_out=../shared/pkg/grpc/orchestrator/ --go-grpc_opt=paths=source_relative orchestrator.proto
 	@echo "Done"
 

packages/orchestrator/go.mod

@@ -98,6 +98,7 @@ require (
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mattn/go-sqlite3 v1.14.27 // indirect
 	github.com/mdlayher/genetlink v1.3.2 // indirect
 	github.com/mdlayher/netlink v1.7.2 // indirect
 	github.com/mdlayher/socket v0.5.1 // indirect

packages/orchestrator/go.sum

@@ -688,6 +688,8 @@ github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D
 github.com/mattn/go-runewidth v0.0.2/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
 github.com/mattn/go-shellwords v1.0.3/go.mod h1:3xCvwCdWdlDJUrvuMn7Wuy9eWs4pE8vqg+NOMyg4B2o=
 github.com/mattn/go-shellwords v1.0.12/go.mod h1:EZzvwXDESEeg03EKmM+RmDnNOPKG4lLtQsUlTZDWQ8Y=
+github.com/mattn/go-sqlite3 v1.14.27 h1:drZCnuvf37yPfs95E5jd9s3XhdVWLal+6BOK6qrv6IU=
+github.com/mattn/go-sqlite3 v1.14.27/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4=
 github.com/mdlayher/genetlink v1.3.2 h1:KdrNKe+CTu+IbZnm/GVUMXSqBBLqcGpRDa0xkQy56gw=

packages/orchestrator/internal/db/db.go

@@ -0,0 +1,132 @@
+package db
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/pkg/database"
+)
+
+type DB struct {
+	client *sql.DB
+	ops    *database.Queries
+}
+
+func New(ctx context.Context, client *sql.DB, schema string) (*DB, error) {
+	db := &DB{client: client, ops: database.New(client)}
+
+	if _, err := db.client.ExecContext(ctx, schema); err != nil {
+		return nil, fmt.Errorf("problem initializing the database: %w", err)
+	}
+
+	if err := db.ops.SetOrchestratorStatusRunning(ctx); err != nil {
+		return nil, err
+	}
+
+	return db, nil
+}
+
+func (db *DB) Close(ctx context.Context) error {
+	return errors.Join(db.ops.SetOrchestratorStatusTerminated(ctx), db.client.Close())
+}
+
+func (db *DB) Status(ctx context.Context) (*database.OrchestratorStatusRow, error) {
+	report, err := db.ops.OrchestratorStatus(ctx)
+	if err != nil {
+		return nil, err
+	}
+	out := report
+
+	return &out, nil
+}
+
+func (db *DB) CreateSandbox(ctx context.Context, params database.CreateSandboxParams) error {
+	return db.WithTx(ctx, func(ctx context.Context, op *database.Queries) error {
+		if _, err := op.IncGlobalVersion(ctx); err != nil {
+			return err
+		}
+
+		if err := op.CreateSandbox(ctx, params); err != nil {
+			return err
+		}
+		return nil
+	})
+}
+
+func (db *DB) UpdateSandboxDeadline(ctx context.Context, id string, deadline time.Time) error {
+	return db.WithTx(ctx, func(ctx context.Context, op *database.Queries) error {
+		if _, err := op.IncGlobalVersion(ctx); err != nil {
+			return err
+		}
+
+		if err := op.UpdateSandboxDeadline(ctx, database.UpdateSandboxDeadlineParams{ID: id, Deadline: deadline}); err != nil {
+			return err
+		}
+
+		return nil
+	})
+}
+
+func (db *DB) SetSandboxTerminated(ctx context.Context, id string, duration time.Duration) error {
+	return db.WithTx(ctx, func(ctx context.Context, op *database.Queries) error {
+		if _, err := op.IncGlobalVersion(ctx); err != nil {
+			return err
+		}
+
+		if err := op.ShutdownSandbox(ctx, database.ShutdownSandboxParams{
+			ID:         id,
+			DurationMs: duration.Milliseconds(),
+			Status:     database.SandboxStatusTerminated,
+		}); err != nil {
+			return err
+		}
+
+		return nil
+	})
+}
+
+func (db *DB) SetSandboxPaused(ctx context.Context, id string, duration time.Duration) error {
+	return db.WithTx(ctx, func(ctx context.Context, op *database.Queries) error {
+		if _, err := op.IncGlobalVersion(ctx); err != nil {
+			return err
+		}
+
+		if err := op.ShutdownSandbox(ctx, database.ShutdownSandboxParams{
+			ID:         id,
+			DurationMs: duration.Milliseconds(),
+			Status:     database.SandboxStatusPaused,
+		}); err != nil {
+			return err
+		}
+
+		return nil
+	})
+}
+
+func (db *DB) WithTx(ctx context.Context, op func(context.Context, *database.Queries) error) (err error) {
+	tx, err := db.client.BeginTx(ctx, nil)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err != nil {
+			err = errors.Join(err, tx.Rollback())
+		}
+	}()
+	defer func() {
+		if p := recover(); p != nil {
+			tx.Rollback()
+			panic(p)
+		}
+	}()
+	defer func() {
+		if err == nil {
+			err = tx.Commit()
+		}
+	}()
+
+	return op(ctx, db.ops.WithTx(tx))
+}

packages/orchestrator/internal/db/queries.sql

@@ -0,0 +1,76 @@
+-- name: GlobalVersion :one
+SELECT version FROM status WHERE id = 1;
+
+-- name: SetOrchestratorStatusRunning :exec
+UPDATE status
+SET
+   version = version + 1,
+   updated_at = current_timestamp,
+   status = 'running'
+WHERE
+   id = 1 AND status = 'initializing';
+
+-- name: SetOrchestratorStatusTerminated :exec
+UPDATE status
+SET
+   version = version + 1,
+   updated_at = current_timestamp,
+   status = 'terminated'
+WHERE
+   id = 1 AND status != 'terminated';
+
+-- name: IncGlobalVersion :one
+UPDATE status
+SET
+   version = version + 1,
+   updated_at = current_timestamp
+WHERE
+   id = 1 AND status != 'terminated'
+RETURNING version;
+
+-- name: CreateSandbox :exec
+INSERT INTO sandboxes(id, status, started_at, deadline, config, global_version)
+VALUES (
+   sqlc.arg('id'),
+   sqlc.arg('status'),
+   sqlc.arg('started_at'),
+   sqlc.arg('deadline'),
+   sqlc.arg('config'),
+   (SELECT version FROM status WHERE status.id = 1)
+);
+
+-- name: ShutdownSandbox :exec
+UPDATE sandboxes
+SET
+  version = version + 1,
+  global_version = (SELECT version FROM status WHERE status.id = 1),
+  updated_at = current_timestamp,
+  duration_ms = sqlc.arg('duration_ms'),
+  status = sqlc.arg('status')
+WHERE
+  sandboxes.id = sqlc.arg('id');
+
+-- name: UpdateSandboxDeadline :exec
+UPDATE sandboxes
+SET
+  version = version + 1,
+  global_version = (SELECT version FROM status WHERE status.id = 1),
+  udpated_at = current_timestamp,
+  deadline = sqlc.arg('deadline'),
+  status = 'running'
+WHERE
+  sandboxes.id = sqlc.arg('id');
+
+-- name: OrchestratorStatus :one
+SELECT
+  status.version AS global_version,
+  (SELECT count(*) FROM sandboxes) AS num_sandboxes,
+  (SELECT count(*) FROM sandboxes WHERE status = 'pending') AS pending_sandboxes,
+  (SELECT count(*) FROM sandboxes WHERE status = 'terminated') AS terminated_sandboxes,
+  (SELECT count(*) FROM sandboxes WHERE status = 'running') AS running_sandboxes,
+  (SELECT min(started_at) FROM sandboxes WHERE status = 'running') AS earliest_running_sandbox_started_at,
+  (SELECT max(updated_at) FROM sandboxes WHERE status = 'running') AS most_recent_running_sandbox_updated_at,
+  status.updated_at,
+  status.status
+FROM
+  status;

packages/orchestrator/internal/db/schema.sql

@@ -0,0 +1,28 @@
+PRAGMA integrity_check;
+PRAGMA foreign_keys = ON;
+
+CREATE TABLE sandboxes (
+  id              TEXT        PRIMARY KEY  NOT NULL,
+  started_at      TIMESTAMP   NOT NULL     DEFAULT current_timestamp,
+  updated_at      TIMESTAMP   NOT NULL     DEFAULT current_timestamp,
+  deadline        TIMESTAMP   NOT NULL,
+  status          TEXT        CHECK( status IN ('pending', 'running', 'paused', 'terminated'))
+                              NOT NULL     DEFAULT 'pending',
+  duration_ms     INTEGER     CHECK( duration_ms >= 0 )
+                              NOT NULL     DEFAULT 0,
+  version         INTEGER     CHECK( version >= 0 )
+                              NOT NULL     DEFAULT 0,
+  global_version  INTEGER     CHECK( global_version >= 0 )
+                              NOT NULL,
+  config          BLOB
+);
+
+CREATE TABLE status (
+  id              INTEGER     PRIMARY KEY  NOT NULL,
+  version         INTEGER     NOT NULL     DEFAULT 0,
+  updated_at      TIMESTAMP   NOT NULL     DEFAULT current_timestamp,
+  status          TEXT        CHECK( status IN ('initializing', 'running', 'draining', 'quarantined ', 'terminated'))
+                              NOT NULL     DEFAULT 'initializing'
+);
+
+INSERT INTO status(id) VALUES(1);

packages/orchestrator/internal/pkg/database/constants.go

@@ -0,0 +1,14 @@
+package database
+
+const (
+	SandboxStatusPending    = "pending"
+	SandboxStatusRunning    = "running"
+	SandboxStatusPaused     = "paused"
+	SandboxStatusTerminated = "terminated"
+
+	OrchestratorStatusInitializing = "initializing"
+	OrchestratorStatusRunning      = "running"
+	OrchestratorStatusDraining     = "draining"
+	OrchestratorStatusTerminated   = "terminated"
+	OrchestratorStatusQuarantined  = "quarantined"
+)