diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2d3b86577..97a2c8d75 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,7 +19,7 @@ jobs: shell: bash steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Install Docker and Prereqs # This uses a workaround for a known issue with docker. See here: https://github.com/actions/virtual-environments/issues/5490#issuecomment-1118328567 run: | @@ -40,11 +40,11 @@ jobs: with: workspaces: "wasm-worker-> target" - name: Install golang - uses: actions/setup-go@v5 + uses: actions/setup-go@v6 with: go-version: '>=1.21' - name: Set up Python 3 - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: '3.12' - name: Install pylint @@ -52,7 +52,7 @@ jobs: - name: Install revive (go linter) run: go install github.com/mgechev/revive@v1.5.1 - name: Install cross - run: cargo install cross + run: cargo install cross --locked - name: Build OpenLambda run: | make ol imgs/ol-wasm wasm-worker wasm-functions native-functions container-proxy @@ -62,12 +62,16 @@ jobs: working-directory: go/common run: go test -v timeout-minutes: 5 - - name: Test Python (Docker) - # not all features are supported by docker yet, so we only run some of the tests - run: sudo env "PATH=$PATH" ./scripts/test.py --worker_type=docker --test_blocklist=max_mem_alloc + - name: Test Go Worker + working-directory: go/worker + run: go test -v ./... + timeout-minutes: 5 - name: Test Python (SOCK) run: sudo env "PATH=$PATH" ./scripts/test.py --worker_type=sock timeout-minutes: 20 + - name: Test Python (Docker) + # not all features are supported by docker yet, so we only run some of the tests + run: sudo env "PATH=$PATH" ./scripts/test.py --worker_type=docker --test_blocklist=max_mem_alloc - name: Boss/Lambda Store Tests run: sudo env "PATH=$PATH" python3 scripts/boss_test.py local timeout-minutes: 5 diff --git a/.github/workflows/pkg.yml b/.github/workflows/pkg.yml index d4fe13723..7d0347dd2 100644 --- a/.github/workflows/pkg.yml +++ b/.github/workflows/pkg.yml @@ -15,7 +15,7 @@ jobs: shell: bash steps: - name: Checkout repository - uses: actions/checkout@v2 + uses: actions/checkout@v5 - name: Install Docker and Prereqs # This uses a workaround for a known issue with docker. See here: https://github.com/actions/virtual-environments/issues/5490#issuecomment-1118328567 run: | @@ -40,11 +40,11 @@ jobs: with: working-directory: wasm-worker - name: Install golang - uses: actions/setup-go@v3 + uses: actions/setup-go@v6 with: go-version: '>=1.21.0' - name: Install cross - run: cargo install cross + run: cargo install cross --locked - name: Build OpenLambda run: | make ol imgs/lambda wasm-worker wasm-functions native-functions container-proxy diff --git a/bin-functions/Cargo.lock b/bin-functions/Cargo.lock index ef6be880b..b42ab06a3 100644 --- a/bin-functions/Cargo.lock +++ b/bin-functions/Cargo.lock @@ -135,9 +135,9 @@ checksum = "c3ac9f8b63eca6fd385229b3675f6cc0dc5c8a5c8a54a59d4f52ffd670d87b0c" [[package]] name = "bytes" -version = "1.10.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f61dac84819c6588b558454b194026eb1f09c293b9036ae9b159e74e73ab6cf9" +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" [[package]] name = "cfg-if" diff --git a/container-proxy/Cargo.lock b/container-proxy/Cargo.lock index 9109f3bc5..85220450f 100644 --- a/container-proxy/Cargo.lock +++ b/container-proxy/Cargo.lock @@ -79,9 +79,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.10.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f61dac84819c6588b558454b194026eb1f09c293b9036ae9b159e74e73ab6cf9" +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" [[package]] name = "cfg-if" diff --git a/docs/worker/README.md b/docs/worker/README.md index 4d8ac7f43..fa1f409f6 100644 --- a/docs/worker/README.md +++ b/docs/worker/README.md @@ -10,6 +10,7 @@ See how to build the code, deploy a worker, and write/run lambdas [here](getting * [dependency management](pypi-packages.md) * [manual cleanup](manual-cleanup.md) * [lambda configuration](lambda-config.md) +* [deploying example applications](apps.md) * [setup development environment](../boss/setup-dev-env.md) * registry config (TODO) * Zygote tree (TODO) @@ -63,7 +64,18 @@ Major cloud offerings (like AWS lambda) offer a variety of lambda triggers, such as HTTP requests, queue messages, cron, DB/S3 triggers, etc. -This has not been a focus (so far) of OpenLambda. The only trigger is -an HTTP request. Thus, all the event code is in the -github.com/open-lambda/open-lambda/ol/worker/server package. Requests -to http(s)://WORKER_ADDR:PORT/run/LAMBDA_NAME invoke lambdas. +OpenLambda currently supports three types of triggers: **HTTP +requests**, **Kafka messages**, and **cron schedules**. The event code +is in the github.com/open-lambda/open-lambda/ol/worker/event package. + +**HTTP triggers:** Requests to http(s)://WORKER_ADDR:PORT/run/LAMBDA_NAME +invoke lambdas directly. + +**Kafka triggers:** Lambdas can be configured to consume from Kafka +topics. The worker runs Kafka consumers that poll for messages and +invoke the corresponding lambda function automatically. See +[kafka-triggers.md](kafka-triggers.md) for details. + +**Cron triggers:** Lambdas can be invoked on a schedule using cron +expressions. The boss component runs a cron scheduler that +automatically invokes lambdas at the configured times. diff --git a/docs/worker/apps.md b/docs/worker/apps.md new file mode 100644 index 000000000..f79265b52 --- /dev/null +++ b/docs/worker/apps.md @@ -0,0 +1,64 @@ +# Deploying Applications + +## Agricultural Forecasting API (FastAPI) + +[ag_forecasting_api](https://github.com/UW-Madison-DSI/ag_forecasting_api) is a FastAPI application that provides crop disease forecasting models for corn and soybean in Wisconsin, developed by University of Wisconsin-Madison plant pathology experts. + +Initialize a worker with the min image: + +```bash +./ol worker init -i ol-min +``` + +Edit `myworker/config.json` to increase memory limit (512MB needed for this app): + +```json +"limits": { + "mem_mb": 512, + ... +} +``` + +Start the worker: + +```bash +./ol worker up -d +``` + +Create `ol.yaml` to configure the app for OpenLambda: + +```yaml +triggers: + http: + - method: "*" +environment: + OL_ENTRY_FILE: app.py + OL_ASGI_ENTRY: app + MEASUREMENTS_CACHE_DIR: /host/tmp/cache + STATIONS_CACHE_FILE: /host/tmp/cache/wisconsin_stations_cache.csv +``` + +Install pip-compile and pin requirements.txt to versions suitable for OpenLambda: + +```bash +./ol admin install examples/pip-compile +curl -X POST -d 'https://raw.githubusercontent.com/tylerharter/ag_forecasting_api/main/requirements.txt' http://localhost:5000/run/pip-compile/url > requirements.txt +``` + +Install and test: + +```bash +./ol admin install -c ol.yaml -r requirements.txt https://github.com/tylerharter/ag_forecasting_api.git + +# simple test +curl http://localhost:5000/run/ag_forecasting_api/ + +# get a forecast for the ALTN station +curl "http://localhost:5000/run/ag_forecasting_api/ag_models_wrappers/wisconet?forecasting_date=2024-07-01&risk_days=1&station_id=ALTN" +``` + +Note, the first request may take minutes because OpenLambda will install all the packages in requirements.txt upon the first call. + +TODO: update ag_forecasting_api URLs from tylerharter fork to UW-Madison-DSI once env option is merged upstream. + +## TODO: add more example apps diff --git a/docs/worker/getting-started.md b/docs/worker/getting-started.md index 117822218..b89b1fcfd 100644 --- a/docs/worker/getting-started.md +++ b/docs/worker/getting-started.md @@ -197,6 +197,22 @@ If you initialized a worker with a specific path (e.g., `./ol worker init -p myw If no `-p` flag is specified, the command will default to the worker running on port 5000 using the default config. +### Installing from a Git Repository + +You can also install lambdas directly from a Git repository (GitHub, GitLab, etc.): + +```bash +./ol admin install https://github.com/open-lambda/hello-lambda-example.git +``` + +This works with both HTTPS and SSH URLs: + +```bash +./ol admin install git@github.com:open-lambda/hello-lambda-example.git +``` + +The function name is derived from the repository name (e.g., `hello-lambda-example`). + ## Invoke Lambda Invoke your lambda with `curl` (the result should be the same as the POST body): diff --git a/docs/worker/kafka-triggers.md b/docs/worker/kafka-triggers.md new file mode 100644 index 000000000..bb2045ec6 --- /dev/null +++ b/docs/worker/kafka-triggers.md @@ -0,0 +1,198 @@ +# Kafka Triggers + +Lambdas can be configured to automatically consume messages from Kafka +topics. When a message arrives, the worker invokes the lambda with the +message payload as the request body. + +## Configuration + +Add a `kafka` section under `triggers` in your lambda's `ol.yaml` +(see [lambda configuration](lambda-config.md) for the full `ol.yaml` +reference): + +```yaml +triggers: + kafka: + - bootstrap_servers: + - "localhost:9092" + topics: + - "my-topic" + auto_offset_reset: "latest" # or "earliest" +``` + +The consumer group ID is automatically set to `lambda-` based on +the lambda name and cannot be overridden. Because each lambda gets its +own group ID, lambdas consume from Kafka independently of one another. +Even if multiple lambdas subscribe to the same topic, each one receives +its own copy of every message, and offset tracking is maintained +separately per lambda. + +## Quick start + +This walkthrough starts a Kafka broker, deploys a lambda with a Kafka +trigger, and publishes a message to verify end-to-end. + +### 1. Start a Kafka broker + +The easiest way to get a single-node broker is with Docker. The +[apache/kafka](https://hub.docker.com/r/apache/kafka) image bundles +KRaft mode so no separate ZooKeeper container is needed: + +```bash +docker run -d --name kafka \ + -p 9092:9092 \ + apache/kafka:latest +``` + +See the [Apache Kafka quickstart](https://kafka.apache.org/quickstart) +for more details. + +### 2. Create a topic + +```bash +docker exec kafka \ + /opt/kafka/bin/kafka-topics.sh --create \ + --topic my-topic \ + --bootstrap-server localhost:9092 +``` + +### 3. Write the lambda + +Create a directory for the lambda with two files: + +**f.py** +```python +def f(event): + print(f"Received: {event}") + return {"status": "ok"} +``` + +**ol.yaml** +```yaml +triggers: + kafka: + - bootstrap_servers: + - "localhost:9092" + topics: + - "my-topic" + auto_offset_reset: "earliest" +``` + +Upload the lambda to the registry. When a lambda with Kafka triggers is +uploaded, the worker automatically starts consumers for the configured +topics — no extra registration step is needed. + +### 4. Publish a test message + +```bash +echo '{"hello":"world"}' | docker exec -i kafka \ + /opt/kafka/bin/kafka-console-producer.sh \ + --topic my-topic \ + --bootstrap-server localhost:9092 +``` + +The worker should pick up the message and invoke your lambda. Check the +worker logs to confirm. + +## How it works + +1. When the worker starts in `lambda` mode, it creates a `KafkaManager` + alongside the `LambdaServer`. +2. When a lambda with Kafka triggers is uploaded, the boss automatically + registers its Kafka consumers on the worker. Consumers can also be + managed manually via the `/kafka/register/` HTTP + endpoint (POST to register, DELETE to unregister). +3. For each trigger entry, the manager creates a `LambdaKafkaConsumer` + backed by a [franz-go](https://github.com/twmb/franz-go) (`kgo`) + client. +4. Each consumer runs a polling loop that fetches messages with a + 1-second timeout. On receiving a message, it builds a synthetic HTTP + POST request and invokes the lambda directly through the + `LambdaManager`. + +## Request format + +When a Kafka message triggers a lambda, the worker builds a synthetic +HTTP POST request with the Kafka message value as the body and the +following headers: + +| Header | Description | +| ------------------- | ---------------------------------------- | +| `Content-Type` | `application/json` | +| `X-Kafka-Topic` | The topic the message was read from. | +| `X-Kafka-Partition` | The partition number. | +| `X-Kafka-Offset` | The message offset within the partition. | +| `X-Kafka-Group-Id` | The consumer group ID. | + +### Accessing Kafka metadata in your handler + +The default handler type (`def f(event)`) only receives the JSON-parsed +request body as a dict. It does **not** have access to HTTP headers, +so the Kafka metadata headers listed above will not be available. + +To access Kafka metadata headers, use a **WSGI** or **ASGI** entry +point (see [lambda configuration](lambda-config.md) for how to +configure these). + +## Example lambdas + +Complete working examples are available in the +[examples/](../../examples/) directory: + +- [kafka-basic](../../examples/kafka-basic/) — Simple `f(event)` handler + that processes the Kafka message body. +- [kafka-metadata](../../examples/kafka-metadata/) — Flask WSGI handler + that accesses Kafka metadata headers (topic, partition, offset, group + ID) alongside the message body. + +### Simple handler (body only) + +The default `f(event)` handler receives the Kafka message body as a +parsed dict, but cannot access headers +([full example](../../examples/kafka-basic/)): + +```python +def f(event): + # event is the JSON-parsed Kafka message value + print(f"Received message: {event}") + return {"status": "ok"} +``` + +### WSGI handler (body + headers) + +A WSGI handler can access Kafka metadata via the `environ` dict. +HTTP headers are available with an `HTTP_` prefix, uppercased, and +with dashes replaced by underscores +([full example](../../examples/kafka-metadata/)): + +```python +from flask import Flask, request + +app = Flask(__name__) + +@app.route("/", methods=["POST"]) +def handle(): + topic = request.headers.get("X-Kafka-Topic", "unknown") + partition = request.headers.get("X-Kafka-Partition", "unknown") + offset = request.headers.get("X-Kafka-Offset", "unknown") + group_id = request.headers.get("X-Kafka-Group-Id", "unknown") + + body = request.get_json() + + print(f"topic={topic} partition={partition} offset={offset} group={group_id}") + print(f"body={body}") + + return {"status": "ok"} +``` + +## Management API + +The worker exposes an HTTP endpoint for managing Kafka consumers at +runtime: + +- **`POST /kafka/register/`** — Reads the lambda's + `ol.yaml` config from the registry and starts consumers for all + configured Kafka triggers. Any existing consumers for that lambda are + cleaned up first. +- **`DELETE /kafka/register/`** — Stops and removes all + Kafka consumers for the given lambda. diff --git a/docs/worker/lambda-config.md b/docs/worker/lambda-config.md index 2ce9aa11b..63462a454 100644 --- a/docs/worker/lambda-config.md +++ b/docs/worker/lambda-config.md @@ -13,12 +13,23 @@ triggers: http: - method: PUT - method: PATCH + kafka: + - bootstrap_servers: + - "localhost:9092" + topics: + - "my-topic" + +environment: + MY_ENV_VAR1: "value1" + MY_ENV_VAR2: "value2" ``` -## 3. Trigger Types -OpenLambda only supports HTTP trigger for now, but future development plans include supporting other trigger types. +## 3. Configuration Options + +### a. Triggers +OpenLambda currently supports HTTP and Kafka triggers. -### a. HTTP Triggers +#### HTTP Triggers Defines which HTTP methods can be used to invoke the lambda. Example: @@ -30,6 +41,84 @@ triggers: ``` In this case, the lambda accepts GET and POST requests. +#### Kafka Triggers +Defines Kafka topics the lambda should consume from. When a message +arrives on a configured topic, the lambda is invoked with the message +as the request body. + +Example: +```yaml +triggers: + kafka: + - bootstrap_servers: + - "localhost:9092" + topics: + - "my-topic" + auto_offset_reset: "latest" +``` + +| Field | Type | Required | Description | +| ------------------- | ---------- | -------- | ------------------------------------------------------------------------------------------- | +| `bootstrap_servers` | `[]string` | Yes | List of Kafka broker addresses. | +| `topics` | `[]string` | Yes | Topics this lambda should consume from. | +| `auto_offset_reset` | `string` | No | Where to start reading if no committed offset exists. `"latest"` (default) or `"earliest"`. | + +A lambda can define multiple Kafka trigger entries. Each entry creates a +separate consumer. For more details on Kafka triggers, including how to access Kafka +metadata headers in your handler, see [kafka-triggers.md](kafka-triggers.md). + +### b. Environment Variables +Defines environment variables that will be available to the lambda function at runtime. + +Example: +```yaml +environment: + MY_ENV_VAR1: "production" + MY_ENV_VAR2: "enabled" +``` + +These variables can be accessed in your lambda code using standard environment variable methods (e.g., `os.environ` in Python). + +**Note:** Environment variables defined in `ol.yaml` are written to a `.env` file in the lambda's directory during execution. If your lambda already has a `.env` file, it will be overwritten with the values from `ol.yaml`. + +### c. Special Environment Variables + +#### OL_ENTRY_FILE +By default, OpenLambda expects Python lambda functions to be defined in a file named `f.py`. You can override this by setting the `OL_ENTRY_FILE` environment variable to specify a different entry file. + +Example: +```yaml +environment: + OL_ENTRY_FILE: "app.py" +``` + +With this configuration: +- OpenLambda will look for `app.py` instead of `f.py` when detecting the Python runtime +- The Python runtime will import the `app` module instead of `f` +- For standard functions, define your handler as `def f(event)` in the specified file +- For Flask/WSGI applications, define your `app` object in the specified file + +This is useful when you want to use conventional naming (e.g., `app.py` for Flask applications) or integrate existing code without renaming files. + +### d. Sandbox Reuse + +#### reuse_sandbox +By default, OpenLambda reuses the same sandbox across multiple invocations of a lambda function to improve performance. In some cases, such as when strict isolation is required or when avoiding state persistence between invocations, it may be desirable to create a fresh sandbox for each invocation. + +This behavior can be controlled using the reuse-sandbox option. + +Example: +```yaml +reuse-sandbox: false +``` + +With this configuration: + +- A new sandbox is created for each lambda invocation +- The sandbox is destroyed after the invocation completes + +If reuse-sandbox is not specified, OpenLambda defaults to reusing sandboxes across invocations. + ## 4. How to Use ### a. Define Configuration Create an `ol.yaml` file inside the lambda function directory with the desired configuration. diff --git a/docs/worker/pypi-packages.md b/docs/worker/pypi-packages.md index 87618e1c0..dc8fa361d 100644 --- a/docs/worker/pypi-packages.md +++ b/docs/worker/pypi-packages.md @@ -45,6 +45,25 @@ six==1.16.0 # via python-dateutil ``` +### Using the pip-compile Lambda + +If you don't have pip-tools installed locally, you can use the +`pip-compile` lambda included in OpenLambda. First, install it: + +```bash +ol admin install ./examples/pip-compile +``` + +Then compile your requirements (from a file or URL): + +```bash +# From a local file +curl -X POST --data-binary @requirements.in http://localhost:5000/run/pip-compile/text > requirements.txt + +# From a URL +curl -X POST -d 'https://example.com/requirements.in' http://localhost:5000/run/pip-compile/url > requirements.txt +``` + ## Try It Start an OpenLambda worker (if not already started). For example, you diff --git a/examples/env-test/f.py b/examples/env-test/f.py new file mode 100644 index 000000000..7fe1f5f4a --- /dev/null +++ b/examples/env-test/f.py @@ -0,0 +1,30 @@ +import os +import json + +def f(event): + """ + Lambda function that demonstrates environment variable usage. + Returns all environment variables that were configured in ol.yaml + """ + + # Get environment variables from config + env_vars = { + "MY_ENV_VAR": os.environ.get("MY_ENV_VAR", "not set"), + "DATABASE_URL": os.environ.get("DATABASE_URL", "not set"), + "DEBUG_MODE": os.environ.get("DEBUG_MODE", "not set"), + "API_KEY": os.environ.get("API_KEY", "not set"), + "CUSTOM_PATH": os.environ.get("CUSTOM_PATH", "not set"), + } + + response = { + "message": "Environment variables test", + "event": event, + "configured_env_vars": env_vars, + "all_env_vars_count": len(os.environ), + } + + # If debug mode is enabled, show all environment variables + if os.environ.get("DEBUG_MODE") == "true": + response["all_env_vars"] = dict(os.environ) + + return response \ No newline at end of file diff --git a/examples/env-test/ol.yaml b/examples/env-test/ol.yaml new file mode 100644 index 000000000..387a52a3a --- /dev/null +++ b/examples/env-test/ol.yaml @@ -0,0 +1,11 @@ +triggers: + http: + - method: GET + - method: POST + +environment: + MY_ENV_VAR: "Hello from environment" + DATABASE_URL: "postgresql://user:pass@localhost/db" + DEBUG_MODE: "true" + API_KEY: "secret-key-789" + CUSTOM_PATH: "/usr/local/bin" \ No newline at end of file diff --git a/examples/fastapi-test/f.py b/examples/fastapi-test/f.py new file mode 100644 index 000000000..44f9163e7 --- /dev/null +++ b/examples/fastapi-test/f.py @@ -0,0 +1,7 @@ +from fastapi import FastAPI + +app = FastAPI() + +@app.get("/") +def hello(): + return {"message": "hello world"} diff --git a/examples/fastapi-test/requirements.in b/examples/fastapi-test/requirements.in new file mode 100644 index 000000000..6b0b9396e --- /dev/null +++ b/examples/fastapi-test/requirements.in @@ -0,0 +1 @@ +fastapi diff --git a/examples/fastapi-test/requirements.txt b/examples/fastapi-test/requirements.txt new file mode 100644 index 000000000..04d1cbda0 --- /dev/null +++ b/examples/fastapi-test/requirements.txt @@ -0,0 +1,12 @@ + +annotated-doc==0.0.4 +annotated-types==0.7.0 +anyio==4.12.1 +exceptiongroup==1.3.1 +fastapi==0.128.0 +idna==3.11 +pydantic-core==2.41.5 +pydantic==2.12.5 +starlette==0.50.0 +typing-extensions==4.15.0 +typing-inspection==0.4.2 diff --git a/examples/flask-entry-test/app.py b/examples/flask-entry-test/app.py new file mode 100644 index 000000000..fa1b6fc60 --- /dev/null +++ b/examples/flask-entry-test/app.py @@ -0,0 +1,14 @@ +from flask import Flask, request, Response + +app = Flask("flask-entry-test") + +@app.route("/") +def index(): + return Response("Hello from app.py!\n", status=200) + +@app.route("/info") +def info(): + return { + "entry_file": "app.py", + "message": "This function uses OL_ENTRY_FILE to specify app.py as the entry point" + } diff --git a/examples/flask-entry-test/ol.yaml b/examples/flask-entry-test/ol.yaml new file mode 100644 index 000000000..fa6be1c13 --- /dev/null +++ b/examples/flask-entry-test/ol.yaml @@ -0,0 +1,7 @@ +triggers: + http: + - method: GET + - method: POST + +environment: + OL_ENTRY_FILE: "app.py" diff --git a/examples/flask-entry-test/requirements.in b/examples/flask-entry-test/requirements.in new file mode 100644 index 000000000..944e2dae9 --- /dev/null +++ b/examples/flask-entry-test/requirements.in @@ -0,0 +1,2 @@ +flask +werkzeug<3.2 diff --git a/examples/flask-entry-test/requirements.txt b/examples/flask-entry-test/requirements.txt new file mode 100644 index 000000000..8148fdc18 --- /dev/null +++ b/examples/flask-entry-test/requirements.txt @@ -0,0 +1,20 @@ +# This file was autogenerated by uv via the following command: +# uv pip compile requirements.in -o requirements.txt +blinker==1.9.0 + # via flask +click==8.1.7 + # via flask +flask==3.1.0 + # via -r requirements.in +itsdangerous==2.2.0 + # via flask +jinja2==3.1.4 + # via flask +markupsafe==3.0.2 + # via + # jinja2 + # werkzeug +werkzeug==3.1.5 + # via + # -r requirements.in + # flask diff --git a/examples/flask-test/requirements.in b/examples/flask-test/requirements.in index b44f14a28..fba9b4e9e 100644 --- a/examples/flask-test/requirements.in +++ b/examples/flask-test/requirements.in @@ -1,2 +1,2 @@ flask==2.3.2 -werkzeug==3.0.3 +werkzeug==3.1.5 diff --git a/examples/flask-test/requirements.txt b/examples/flask-test/requirements.txt index 96159fc91..bb7ff3d8b 100644 --- a/examples/flask-test/requirements.txt +++ b/examples/flask-test/requirements.txt @@ -18,7 +18,7 @@ markupsafe==2.1.3 # via # jinja2 # werkzeug -werkzeug==3.0.3 +werkzeug==3.1.5 # via # -r requirements.in # flask diff --git a/examples/kafka-basic/f.py b/examples/kafka-basic/f.py new file mode 100644 index 000000000..5009c6e03 --- /dev/null +++ b/examples/kafka-basic/f.py @@ -0,0 +1,4 @@ +def f(event): + # event is the JSON-parsed Kafka message value + print(f"Received message: {event}") + return {"status": "ok"} diff --git a/examples/kafka-basic/ol.yaml b/examples/kafka-basic/ol.yaml new file mode 100644 index 000000000..4e00a5664 --- /dev/null +++ b/examples/kafka-basic/ol.yaml @@ -0,0 +1,9 @@ +triggers: + http: + - method: POST + kafka: + - bootstrap_servers: + - "localhost:9092" + topics: + - "my-topic" + auto_offset_reset: "latest" diff --git a/examples/kafka-db-sum/f.py b/examples/kafka-db-sum/f.py new file mode 100644 index 000000000..957b9a850 --- /dev/null +++ b/examples/kafka-db-sum/f.py @@ -0,0 +1,198 @@ +from flask import Flask, request, make_response, jsonify +import os +import random +import psycopg2 + +app = Flask(__name__) + +DATABASE_URL = os.environ.get("DATABASE_URL") + +# Probability (0.0–1.0) that a transaction will fail between UPDATE and COMMIT. +# Set to 0 for normal operation; raise to stress-test seek-back recovery. +FAIL_PROBABILITY = float(os.environ.get("FAIL_PROBABILITY", "0")) +_db_initialized = False + + +def get_db(): + return psycopg2.connect(DATABASE_URL) + + +def ensure_db(): + """Create the running_sum table if it doesn't exist (runs once per sandbox).""" + global _db_initialized + if _db_initialized: + return + conn = get_db() + try: + with conn.cursor() as cur: + cur.execute( + """ + CREATE TABLE IF NOT EXISTS running_sum ( + id INTEGER PRIMARY KEY DEFAULT 1, + total BIGINT NOT NULL DEFAULT 0, + last_offset BIGINT NOT NULL DEFAULT -1, + message_count BIGINT NOT NULL DEFAULT 0, + CHECK (id = 1) + ) + """ + ) + cur.execute( + """ + INSERT INTO running_sum (id, total, last_offset, message_count) + VALUES (1, 0, -1, 0) + ON CONFLICT (id) DO NOTHING + """ + ) + conn.commit() + _db_initialized = True + finally: + conn.close() + + +@app.route("/reset", methods=["POST"]) +def reset(): + """Reset running_sum to zero so the demo can be re-run cleanly.""" + ensure_db() + conn = get_db() + try: + with conn.cursor() as cur: + cur.execute( + "UPDATE running_sum SET total = 0, last_offset = -1, " + "message_count = 0 WHERE id = 1" + ) + conn.commit() + return jsonify({"status": "reset"}) + finally: + conn.close() + + +@app.route("/", methods=["GET", "POST"]) +def handle(): + ensure_db() + + # GET — return current state (useful for checking progress via HTTP) + if request.method == "GET": + conn = get_db() + try: + with conn.cursor() as cur: + cur.execute( + "SELECT total, last_offset, message_count " + "FROM running_sum WHERE id = 1" + ) + row = cur.fetchone() + if row: + return jsonify( + { + "running_sum": row[0], + "last_offset": row[1], + "message_count": row[2], + } + ) + return jsonify( + {"running_sum": 0, "last_offset": -1, "message_count": 0} + ) + finally: + conn.close() + + # POST — process a Kafka message containing a number + offset = int(request.headers.get("X-Kafka-Offset", "-1")) + topic = request.headers.get("X-Kafka-Topic", "unknown") + partition = request.headers.get("X-Kafka-Partition", "unknown") + + body = request.get_json(silent=True) + + # Accept {"number": N} + if isinstance(body, dict): + number = body.get("number", 0) + + conn = None + try: + conn = get_db() + with conn.cursor() as cur: + # Lock the row and read last processed offset + cur.execute( + "SELECT last_offset FROM running_sum WHERE id = 1 FOR UPDATE" + ) + row = cur.fetchone() + last_offset = row[0] if row else -1 + + # Idempotency: skip if this offset was already processed. + # This prevents double-counting after a seek-back replays + # messages that were already committed. + if offset <= last_offset: + conn.rollback() + print(f"[skip] offset={offset} already processed (last={last_offset})") + return jsonify( + { + "status": "skipped", + "reason": "already processed", + "offset": offset, + "last_offset": last_offset, + } + ) + + # Atomically add number to running sum and advance the offset + cur.execute( + """ + UPDATE running_sum + SET total = total + %s, + last_offset = %s, + message_count = message_count + 1 + WHERE id = 1 + """, + (number, offset), + ) + + # --- Fault injection ------------------------------------------------ + # Simulate a crash between UPDATE and COMMIT. + if FAIL_PROBABILITY > 0 and random.random() < FAIL_PROBABILITY: + raise Exception( + f"Simulated DB failure at offset {offset} " + f"(FAIL_PROBABILITY={FAIL_PROBABILITY})" + ) + # -------------------------------------------------------------------- + + conn.commit() + + # Read back the new state for the response + cur.execute( + "SELECT total, last_offset, message_count " + "FROM running_sum WHERE id = 1" + ) + total, last_off, count = cur.fetchone() + + print(f"[ok] offset={offset} number={number} sum={total} count={count}") + return jsonify( + { + "status": "ok", + "offset": offset, + "number_added": number, + "running_sum": total, + "message_count": count, + } + ) + + except Exception as e: + print(f"[error] offset={offset} error={e}") + if conn: + try: + conn.rollback() + except Exception: + pass + + # Tell OL's Kafka consumer to seek back to this offset and retry. + # The consumer's LRU cache will serve the replay without re-fetching + # from Kafka, and the idempotency check above prevents double-counting + # for any offsets that were already committed before the failure. + resp = make_response( + jsonify({"status": "error", "offset": offset, "error": str(e)}), 500 + ) + resp.headers["X-Kafka-Seek-Offset"] = str(offset) + return resp + + finally: + if conn: + try: + conn.close() + except Exception: + pass diff --git a/examples/kafka-db-sum/instructions.md b/examples/kafka-db-sum/instructions.md new file mode 100644 index 000000000..e9b0872e1 --- /dev/null +++ b/examples/kafka-db-sum/instructions.md @@ -0,0 +1,106 @@ +# kafka-db-sum: Testing Instructions + +## Prerequisites + +- OpenLambda built (`make ol imgs/ol-min`) +- Docker installed + +## 1. Start PostgreSQL + +```bash +docker run -d --name ol-pg \ + --network host \ + -e POSTGRES_USER=ol \ + -e POSTGRES_PASSWORD=ol \ + -e POSTGRES_DB=ol_demo \ + postgres:16 +``` + +## 2. Start Kafka + +```bash +docker run -d --name kafka \ + -p 9092:9092 \ + apache/kafka:latest +``` + +## 3. Create the `numbers` topic + +```bash +docker exec kafka /opt/kafka/bin/kafka-topics.sh --create \ + --topic numbers \ + --bootstrap-server localhost:9092 +``` + +## 4. Initialize and start the OL worker + +From the repository root: + +```bash +sudo -A ./ol worker init -p ../default-ol -i ol-min +sudo -A ./ol worker up -p ../default-ol +``` + +Run `worker up` in a separate terminal, or add `-d` for detached mode. +The worker listens on `localhost:5000` by default. + +## 5. Install the lambda + +From the repository root: + +```bash +./ol admin install examples/kafka-db-sum/ +``` + +## 6. Register the Kafka consumer + +A standalone worker does not auto-register Kafka triggers on upload. +Register manually: + +```bash +curl -X POST localhost:5000/kafka/register/kafka-db-sum +``` + +## 7. Send test messages + +Python producer script (requires `pip install kafka-python`): + +```bash +python examples/kafka-db-sum/produce.py 100 +``` + +## 8. Check results + +```bash +curl localhost:5000/run/kafka-db-sum/ +``` + +Expected output (sum of 1..100 = 5050): + +```json +{ "last_offset": 99, "message_count": 100, "running_sum": 5050 } +``` + +## 9. Reset and re-run + +```bash +curl -X POST localhost:5000/run/kafka-db-sum/reset +``` + +Then send a fresh batch (step 7) and verify again. + +## Configuration + +In `ol.yaml`: + +| Variable | Default | Description | +| ------------------ | ------------------------------------------- | ------------------------------------------------------------------- | +| `DATABASE_URL` | `postgresql://ol:ol@127.0.0.1:5432/ol_demo` | PostgreSQL connection string | +| `FAIL_PROBABILITY` | `0` | Chance (0.0-1.0) of simulated failure. Use `0.3` to test seek-back. | + +## Cleanup + +```bash +sudo -A ./ol worker down -p default-ol +docker rm -f kafka ol-pg +``` diff --git a/examples/kafka-db-sum/ol.yaml b/examples/kafka-db-sum/ol.yaml new file mode 100644 index 000000000..28fd06747 --- /dev/null +++ b/examples/kafka-db-sum/ol.yaml @@ -0,0 +1,16 @@ +triggers: + http: + - method: GET + - method: POST + kafka: + - bootstrap_servers: + - "localhost:9092" + topics: + - "numbers" + auto_offset_reset: "earliest" + +environment: + DATABASE_URL: "postgresql://ol:ol@127.0.0.1:5432/ol_demo" + # Probability (0.0-1.0) of simulated DB failure between UPDATE and COMMIT. + # Set to "0" for normal operation. Try "0.3" to see seek-back recovery in action. + FAIL_PROBABILITY: "0.3" diff --git a/examples/kafka-db-sum/produce.py b/examples/kafka-db-sum/produce.py new file mode 100644 index 000000000..c6a3671ec --- /dev/null +++ b/examples/kafka-db-sum/produce.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +""" +Send numbered messages to the 'numbers' Kafka topic. + +Usage: + python produce.py # send numbers 1..10 + python produce.py 100 # send numbers 1..100 + python produce.py 50 0.5 # send 1..50 with 0.5s delay between each +""" + +import json +import sys +import time + +from kafka import KafkaProducer + +BROKER = "localhost:9092" +TOPIC = "numbers" + + +def main(): + count = int(sys.argv[1]) if len(sys.argv) > 1 else 10 + delay = float(sys.argv[2]) if len(sys.argv) > 2 else 0.1 + + producer = KafkaProducer( + bootstrap_servers=BROKER, + value_serializer=lambda v: json.dumps(v).encode("utf-8"), + ) + + expected_sum = 0 + print(f"Sending numbers 1..{count} to topic '{TOPIC}'...") + for i in range(1, count + 1): + producer.send(TOPIC, {"number": i}) + expected_sum += i + print(f" sent {i}") + if delay: + time.sleep(delay) + + producer.flush() + producer.close() + print(f"\nDone. Expected sum = {expected_sum}") + + +if __name__ == "__main__": + main() diff --git a/examples/kafka-db-sum/requirements.in b/examples/kafka-db-sum/requirements.in new file mode 100644 index 000000000..cefe02572 --- /dev/null +++ b/examples/kafka-db-sum/requirements.in @@ -0,0 +1,3 @@ +flask==2.3.2 +werkzeug==3.0.3 +psycopg2-binary==2.9.9 diff --git a/examples/kafka-db-sum/requirements.txt b/examples/kafka-db-sum/requirements.txt new file mode 100644 index 000000000..869c94ee1 --- /dev/null +++ b/examples/kafka-db-sum/requirements.txt @@ -0,0 +1,26 @@ +# +# This file is autogenerated by pip-compile with Python 3.13 +# by the following command: +# +# pip-compile requirements.in +# +blinker==1.6.2 + # via flask +click==8.1.7 + # via flask +flask==2.3.2 + # via -r requirements.in +itsdangerous==2.1.2 + # via flask +jinja2==3.1.4 + # via flask +markupsafe==2.1.3 + # via + # jinja2 + # werkzeug +psycopg2-binary==2.9.9 + # via -r requirements.in +werkzeug==3.0.3 + # via + # -r requirements.in + # flask diff --git a/examples/kafka-metadata/f.py b/examples/kafka-metadata/f.py new file mode 100644 index 000000000..d92f6f860 --- /dev/null +++ b/examples/kafka-metadata/f.py @@ -0,0 +1,17 @@ +from flask import Flask, request + +app = Flask(__name__) + +@app.route("/", methods=["POST"]) +def handle(): + topic = request.headers.get("X-Kafka-Topic", "unknown") + partition = request.headers.get("X-Kafka-Partition", "unknown") + offset = request.headers.get("X-Kafka-Offset", "unknown") + group_id = request.headers.get("X-Kafka-Group-Id", "unknown") + + body = request.get_json() + + print(f"topic={topic} partition={partition} offset={offset} group={group_id}") + print(f"body={body}") + + return {"status": "ok"} diff --git a/examples/kafka-metadata/ol.yaml b/examples/kafka-metadata/ol.yaml new file mode 100644 index 000000000..319e8bb6f --- /dev/null +++ b/examples/kafka-metadata/ol.yaml @@ -0,0 +1,10 @@ +triggers: + http: + - method: POST + kafka: + - bootstrap_servers: + - "localhost:9092" + topics: + - "my-topic" + auto_offset_reset: "earliest" + diff --git a/examples/kafka-metadata/requirements.in b/examples/kafka-metadata/requirements.in new file mode 100644 index 000000000..b44f14a28 --- /dev/null +++ b/examples/kafka-metadata/requirements.in @@ -0,0 +1,2 @@ +flask==2.3.2 +werkzeug==3.0.3 diff --git a/examples/kafka-metadata/requirements.txt b/examples/kafka-metadata/requirements.txt new file mode 100644 index 000000000..96159fc91 --- /dev/null +++ b/examples/kafka-metadata/requirements.txt @@ -0,0 +1,24 @@ +# +# This file is autogenerated by pip-compile with Python 3.13 +# by the following command: +# +# pip-compile requirements.in +# +blinker==1.6.2 + # via flask +click==8.1.7 + # via flask +flask==2.3.2 + # via -r requirements.in +itsdangerous==2.1.2 + # via flask +jinja2==3.1.4 + # via flask +markupsafe==2.1.3 + # via + # jinja2 + # werkzeug +werkzeug==3.0.3 + # via + # -r requirements.in + # flask diff --git a/examples/lambda-config-test/requirements.in b/examples/lambda-config-test/requirements.in index b44f14a28..fba9b4e9e 100644 --- a/examples/lambda-config-test/requirements.in +++ b/examples/lambda-config-test/requirements.in @@ -1,2 +1,2 @@ flask==2.3.2 -werkzeug==3.0.3 +werkzeug==3.1.5 diff --git a/examples/lambda-config-test/requirements.txt b/examples/lambda-config-test/requirements.txt index 96159fc91..bb7ff3d8b 100644 --- a/examples/lambda-config-test/requirements.txt +++ b/examples/lambda-config-test/requirements.txt @@ -18,7 +18,7 @@ markupsafe==2.1.3 # via # jinja2 # werkzeug -werkzeug==3.0.3 +werkzeug==3.1.5 # via # -r requirements.in # flask diff --git a/examples/pip-compile/f.py b/examples/pip-compile/f.py new file mode 100644 index 000000000..fa0ecd3c8 --- /dev/null +++ b/examples/pip-compile/f.py @@ -0,0 +1,99 @@ +import os +import urllib.request +import urllib.error + +SCRATCH_DIR = "/host/tmp" + +# Set cache/temp directories to writable location BEFORE importing pip-tools +os.environ["HOME"] = SCRATCH_DIR +os.environ["TMPDIR"] = SCRATCH_DIR +os.environ["XDG_CACHE_HOME"] = SCRATCH_DIR +os.environ["PIP_CACHE_DIR"] = SCRATCH_DIR + +from flask import Flask, request, Response +from piptools.scripts.compile import cli +from click.testing import CliRunner + +app = Flask(__name__) + + +def do_compile(requirements_in, quiet=True): + """Compile requirements.in content to requirements.txt.""" + if not requirements_in: + return Response("No requirements provided", status=400, mimetype="text/plain") + + in_path = os.path.join(SCRATCH_DIR, "requirements.in") + out_path = os.path.join(SCRATCH_DIR, "requirements.txt") + + with open(in_path, "w") as f: + f.write(requirements_in) + + args = [ + "--output-file", out_path, + "--pip-args", "--only-binary=:all:", + ] + if quiet: + args.extend(["--no-header", "--no-annotate"]) + args.append(in_path) + + runner = CliRunner() + result = runner.invoke(cli, args) + + if result.exit_code != 0: + return Response( + result.output or str(result.exception), + status=400, + mimetype="text/plain" + ) + + with open(out_path, "r") as f: + lines = [l for l in f if not l.startswith("--")] + return Response("".join(lines), mimetype="text/plain") + + +@app.route("/", methods=["GET"]) +def docs(): + """Return documentation with curl examples.""" + return Response("""pip-compile Lambda Service +========================== + +Compiles requirements.in files to pinned requirements.txt using pip-compile. + +Endpoints +--------- + +POST /text + Pass requirements.in content directly in the request body. + + curl -X POST -d $'flask>=2.0\\nrequests' http://localhost:5000/run/pip-compile/text + +POST /url + Pass a URL to fetch requirements.in from. + + curl -X POST -d 'https://example.com/requirements.in' http://localhost:5000/run/pip-compile/url +""", mimetype="text/plain") + + +@app.route("/text", methods=["POST"]) +def compile_from_text(): + """Compile requirements.in from POST body text.""" + quiet = request.args.get("quiet", "1") == "1" + return do_compile(request.get_data(as_text=True), quiet=quiet) + + +@app.route("/url", methods=["POST"]) +def compile_from_url(): + """Fetch requirements.in from a URL and compile it.""" + url = request.get_data(as_text=True).strip() + quiet = request.args.get("quiet", "1") == "1" + + if not url: + return Response("No URL provided", status=400, mimetype="text/plain") + + try: + with urllib.request.urlopen(url, timeout=30) as response: + requirements_in = response.read().decode('utf-8') + except urllib.error.URLError as e: + return Response(f"Failed to fetch URL: {e}", status=400, mimetype="text/plain") + + return do_compile(requirements_in, quiet=quiet) diff --git a/examples/pip-compile/requirements.in b/examples/pip-compile/requirements.in new file mode 100644 index 000000000..a7c1f43db --- /dev/null +++ b/examples/pip-compile/requirements.in @@ -0,0 +1,2 @@ +flask +pip-tools==5.5.0 diff --git a/examples/pip-compile/requirements.txt b/examples/pip-compile/requirements.txt new file mode 100644 index 000000000..8cb1e9143 --- /dev/null +++ b/examples/pip-compile/requirements.txt @@ -0,0 +1,24 @@ +# +# This file is autogenerated by pip-compile +# +blinker==1.9.0 + # via flask +click==8.3.1 + # via + # flask + # pip-tools +flask==3.1.2 + # via -r requirements.in +itsdangerous==2.2.0 + # via flask +jinja2==3.1.6 + # via flask +markupsafe==3.0.3 + # via + # flask + # jinja2 + # werkzeug +pip-tools==5.5.0 + # via -r requirements.in +werkzeug==3.1.5 + # via flask diff --git a/examples/server/echo.py b/examples/server/echo.py deleted file mode 100644 index f209c6344..000000000 --- a/examples/server/echo.py +++ /dev/null @@ -1,2 +0,0 @@ -def f(event): - return event diff --git a/examples/server/hello.py b/examples/server/hello.py deleted file mode 100644 index f048f4749..000000000 --- a/examples/server/hello.py +++ /dev/null @@ -1,2 +0,0 @@ -def f(event): - return 'hello' diff --git a/examples/wsgi-entry-test/main.py b/examples/wsgi-entry-test/main.py new file mode 100644 index 000000000..e2609ff27 --- /dev/null +++ b/examples/wsgi-entry-test/main.py @@ -0,0 +1,16 @@ +from flask import Flask, Response + +# Intentionally NOT named "app" to test OL_WSGI_ENTRY +my_wsgi_app = Flask("wsgi-entry-test") + +@my_wsgi_app.route("/") +def index(): + return Response("Hello from my_wsgi_app!\n", status=200) + +@my_wsgi_app.route("/info") +def info(): + return { + "entry_file": "main.py", + "entry_point": "my_wsgi_app", + "message": "This tests OL_WSGI_ENTRY with a non-standard name" + } diff --git a/examples/wsgi-entry-test/ol.yaml b/examples/wsgi-entry-test/ol.yaml new file mode 100644 index 000000000..f6172d6c0 --- /dev/null +++ b/examples/wsgi-entry-test/ol.yaml @@ -0,0 +1,8 @@ +triggers: + http: + - method: GET + - method: POST + +environment: + OL_ENTRY_FILE: "main.py" + OL_WSGI_ENTRY: "my_wsgi_app" diff --git a/examples/wsgi-entry-test/requirements.in b/examples/wsgi-entry-test/requirements.in new file mode 100644 index 000000000..fba9b4e9e --- /dev/null +++ b/examples/wsgi-entry-test/requirements.in @@ -0,0 +1,2 @@ +flask==2.3.2 +werkzeug==3.1.5 diff --git a/examples/wsgi-entry-test/requirements.txt b/examples/wsgi-entry-test/requirements.txt new file mode 100644 index 000000000..71333ecbc --- /dev/null +++ b/examples/wsgi-entry-test/requirements.txt @@ -0,0 +1,24 @@ +# +# This file is autogenerated by pip-compile with Python 3.13 +# by the following command: +# +# pip-compile requirements.in +# +blinker==1.6.2 + # via flask +click==8.1.7 + # via flask +flask==2.3.2 + # via -r requirements.in +itsdangerous==2.1.2 + # via flask +jinja2==3.1.6 + # via flask +markupsafe==2.1.3 + # via + # jinja2 + # werkzeug +werkzeug==3.1.5 + # via + # -r requirements.in + # flask diff --git a/examples/wsgi-post-echo/f.py b/examples/wsgi-post-echo/f.py new file mode 100644 index 000000000..e2bc9a6e1 --- /dev/null +++ b/examples/wsgi-post-echo/f.py @@ -0,0 +1,12 @@ +from flask import Flask, request, Response + +app = Flask(__name__) + + +@app.route("/", methods=["GET", "POST", "PUT"]) +def echo(): + """Echo back the POST body.""" + return Response( + request.get_data(as_text=True), + mimetype=request.content_type or "text/plain" + ) diff --git a/examples/wsgi-post-echo/requirements.in b/examples/wsgi-post-echo/requirements.in new file mode 100644 index 000000000..7e1060246 --- /dev/null +++ b/examples/wsgi-post-echo/requirements.in @@ -0,0 +1 @@ +flask diff --git a/examples/wsgi-post-echo/requirements.txt b/examples/wsgi-post-echo/requirements.txt new file mode 100644 index 000000000..6d50b78c9 --- /dev/null +++ b/examples/wsgi-post-echo/requirements.txt @@ -0,0 +1,22 @@ +# +# This file is autogenerated by pip-compile with Python 3.13 +# by the following command: +# +# pip-compile requirements.in +# +blinker==1.6.2 + # via flask +click==8.1.7 + # via flask +flask==2.3.2 + # via -r requirements.in +itsdangerous==2.1.2 + # via flask +jinja2==3.1.6 + # via flask +markupsafe==2.1.3 + # via + # jinja2 + # werkzeug +werkzeug==3.1.5 + # via flask diff --git a/examples/wsgi-test/requirements.in b/examples/wsgi-test/requirements.in index b44f14a28..fba9b4e9e 100644 --- a/examples/wsgi-test/requirements.in +++ b/examples/wsgi-test/requirements.in @@ -1,2 +1,2 @@ flask==2.3.2 -werkzeug==3.0.3 +werkzeug==3.1.5 diff --git a/examples/wsgi-test/requirements.txt b/examples/wsgi-test/requirements.txt index 96159fc91..bb7ff3d8b 100644 --- a/examples/wsgi-test/requirements.txt +++ b/examples/wsgi-test/requirements.txt @@ -18,7 +18,7 @@ markupsafe==2.1.3 # via # jinja2 # werkzeug -werkzeug==3.0.3 +werkzeug==3.1.5 # via # -r requirements.in # flask diff --git a/go/admin/commands.go b/go/admin/commands.go index c4c87189a..e6b6c72f7 100644 --- a/go/admin/commands.go +++ b/go/admin/commands.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "os" + "os/exec" "path/filepath" "strings" "time" @@ -31,13 +32,43 @@ func checkStatus(port string) error { defer resp.Body.Close() if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) + body, readErr := io.ReadAll(resp.Body) + if readErr != nil { + return fmt.Errorf("boss/worker returned status %d (failed to read response body: %v)", resp.StatusCode, readErr) + } return fmt.Errorf("boss/worker returned status %d: %s", resp.StatusCode, string(body)) } return nil } +const installUsage = "ol admin install [-c ] [-r ] [-n ] [boss | -p ] " + +// isGitURL returns true if the path looks like a git repository URL +func isGitURL(path string) bool { + if !strings.HasSuffix(path, ".git") { + return false + } + return strings.HasPrefix(path, "https://") || strings.HasPrefix(path, "git@") +} + +// cloneGitRepo clones a git repository to a temporary directory +func cloneGitRepo(gitURL string) (string, error) { + tmpDir, err := os.MkdirTemp("", "ol-install-*") + if err != nil { + return "", fmt.Errorf("failed to create temp directory: %v", err) + } + + cmd := exec.Command("git", "clone", "--depth", "1", gitURL, tmpDir) + output, err := cmd.CombinedOutput() + if err != nil { + os.RemoveAll(tmpDir) + return "", fmt.Errorf("git clone failed: %v\n%s", err, string(output)) + } + + return tmpDir, nil +} + func adminInstall(ctx *cli.Context) error { args := ctx.Args().Slice() var installTarget string @@ -46,7 +77,7 @@ func adminInstall(ctx *cli.Context) error { workerPath := ctx.String("path") if len(args) == 0 { - return fmt.Errorf("usage: ol admin install [boss | -p ] ") + return fmt.Errorf("usage: %s", installUsage) } if len(args) == 1 { funcDir = args[0] @@ -59,7 +90,7 @@ func adminInstall(ctx *cli.Context) error { return fmt.Errorf("cannot use both 'boss' and '-p' flags together") } } else { - return fmt.Errorf("usage: ol admin install [boss | -p ] ") + return fmt.Errorf("usage: %s", installUsage) } var portToUploadLambda string @@ -96,15 +127,58 @@ func adminInstall(ctx *cli.Context) error { portToUploadLambda = common.Conf.Worker_port } - funcDir = strings.TrimSuffix(funcDir, "/") + var funcName string + var tmpDir string - funcName := filepath.Base(funcDir) + if isGitURL(funcDir) { + funcName = strings.TrimSuffix(filepath.Base(funcDir), ".git") + clonedDir, err := cloneGitRepo(funcDir) + if err != nil { + return err + } + tmpDir = clonedDir + funcDir = clonedDir + } else { + funcDir = strings.TrimSuffix(funcDir, "/") + funcName = filepath.Base(funcDir) + if _, err := os.Stat(funcDir); os.IsNotExist(err) { + return fmt.Errorf("directory %s does not exist", funcDir) + } + } - if _, err := os.Stat(funcDir); os.IsNotExist(err) { - return fmt.Errorf("directory %s does not exist", funcDir) + // Override function name if specified + if name := ctx.String("name"); name != "" { + funcName = name } - tarData, err := createTarGz(funcDir) + // Build overrides map + overrides := make(map[string]string) + addOverride := func(flagName, targetFile string) error { + path := ctx.String(flagName) + if path == "" { + return nil + } + if _, err := os.Stat(path); os.IsNotExist(err) { + return fmt.Errorf("%s file %s does not exist", flagName, path) + } + if _, err := os.Stat(filepath.Join(funcDir, targetFile)); err == nil { + fmt.Printf("Warning: overriding existing %s in source with %s\n", targetFile, path) + } + overrides[targetFile] = path + return nil + } + + if err := addOverride("config", "ol.yaml"); err != nil { + return err + } + if err := addOverride("requirements", "requirements.txt"); err != nil { + return err + } + + tarData, err := createTarGz(funcDir, overrides) + if tmpDir != "" { + os.RemoveAll(tmpDir) + } if err != nil { return fmt.Errorf("failed to create tar.gz: %v", err) } @@ -117,17 +191,34 @@ func adminInstall(ctx *cli.Context) error { return nil } -func createTarGz(funcDir string) ([]byte, error) { +func createTarGz(funcDir string, overrides map[string]string) ([]byte, error) { var buf bytes.Buffer gzWriter := gzip.NewWriter(&buf) tarWriter := tar.NewWriter(gzWriter) - fpyPath := filepath.Join(funcDir, "f.py") - if _, err := os.Stat(fpyPath); os.IsNotExist(err) { - return nil, fmt.Errorf("required file f.py not found in %s", funcDir) + // Determine the Python entry file (default to f.py, or use OL_ENTRY_FILE from ol.yaml) + // Check override config first, then fall back to source config + pythonEntryFile := "f.py" + configDir := funcDir + if configOverride, ok := overrides["ol.yaml"]; ok { + configDir = filepath.Dir(configOverride) + } + lambdaConfig, err := common.LoadLambdaConfig(configDir) + if err != nil { + return nil, fmt.Errorf("failed to parse config in %s: %v", configDir, err) + } + if lambdaConfig.Environment != nil { + if entryFile, ok := lambdaConfig.Environment["OL_ENTRY_FILE"]; ok { + pythonEntryFile = entryFile + } + } + + entryPath := filepath.Join(funcDir, pythonEntryFile) + if _, err := os.Stat(entryPath); os.IsNotExist(err) { + return nil, fmt.Errorf("required file %s not found in %s", pythonEntryFile, funcDir) } - err := filepath.Walk(funcDir, func(path string, info os.FileInfo, err error) error { + err = filepath.Walk(funcDir, func(path string, info os.FileInfo, err error) error { if err != nil { return fmt.Errorf("walk error: %v", err) } @@ -136,11 +227,20 @@ func createTarGz(funcDir string) ([]byte, error) { return nil } + if !info.Mode().IsRegular() { + return fmt.Errorf("cannot archive non-regular file %q (mode: %s)", path, info.Mode().String()) + } + relPath, err := filepath.Rel(funcDir, path) if err != nil { return fmt.Errorf("unable to compute relative path: %v", err) } + // Skip files that will be overridden + if _, ok := overrides[relPath]; ok { + return nil + } + header, err := tar.FileInfoHeader(info, "") if err != nil { return fmt.Errorf("unable to create header: %v", err) @@ -171,6 +271,38 @@ func createTarGz(funcDir string) ([]byte, error) { return nil, err } + // Add override files + for relPath, localPath := range overrides { + info, err := os.Stat(localPath) + if err != nil { + return nil, fmt.Errorf("unable to stat override file %s: %v", localPath, err) + } + + header, err := tar.FileInfoHeader(info, "") + if err != nil { + return nil, fmt.Errorf("unable to create header for override %s: %v", relPath, err) + } + header.Name = relPath + + if err := tarWriter.WriteHeader(header); err != nil { + return nil, fmt.Errorf("failed to write header for override %s: %v", relPath, err) + } + + file, err := os.Open(localPath) + if err != nil { + return nil, fmt.Errorf("unable to open override file %s: %v", localPath, err) + } + + if _, err := io.Copy(tarWriter, file); err != nil { + file.Close() + return nil, fmt.Errorf("error copying override file %s: %v", localPath, err) + } + + if err := file.Close(); err != nil { + return nil, fmt.Errorf("error closing override file %s: %v", localPath, err) + } + } + if err := tarWriter.Close(); err != nil { return nil, fmt.Errorf("failed to close tar writer: %v", err) } @@ -201,7 +333,10 @@ func uploadToLambdaStore(funcName string, tarData []byte, port string) error { defer resp.Body.Close() if resp.StatusCode != http.StatusCreated { - body, _ := io.ReadAll(resp.Body) + body, readErr := io.ReadAll(resp.Body) + if readErr != nil { + return fmt.Errorf("upload failed with status %d (failed to read response body: %v)", resp.StatusCode, readErr) + } return fmt.Errorf("upload failed with status %d: %s", resp.StatusCode, string(body)) } @@ -212,8 +347,8 @@ func AdminCommands() []*cli.Command { return []*cli.Command{ { Name: "install", - Usage: "Install a lambda function from directory", - UsageText: "ol admin install [boss | -p ] ", + Usage: "Install a lambda function from directory or git repo", + UsageText: installUsage, Action: adminInstall, Flags: []cli.Flag{ &cli.StringFlag{ @@ -221,6 +356,21 @@ func AdminCommands() []*cli.Command { Aliases: []string{"p"}, Usage: "Worker directory path (e.g., -p myworker)", }, + &cli.StringFlag{ + Name: "config", + Aliases: []string{"c"}, + Usage: "Path to ol.yaml config file to include (overrides existing ol.yaml in source)", + }, + &cli.StringFlag{ + Name: "requirements", + Aliases: []string{"r"}, + Usage: "Path to requirements.txt file to include (overrides existing requirements.txt in source)", + }, + &cli.StringFlag{ + Name: "name", + Aliases: []string{"n"}, + Usage: "Lambda function name (defaults to directory or repo name)", + }, }, }, } diff --git a/go/boss/cloudvm/local_worker.go b/go/boss/cloudvm/local_worker.go index c373f73bf..8ff86c74f 100644 --- a/go/boss/cloudvm/local_worker.go +++ b/go/boss/cloudvm/local_worker.go @@ -31,7 +31,7 @@ func NewLocalWorkerPool() *WorkerPool { // Get the worker config struct defaultTemplateConfig, err := common.GetDefaultWorkerConfig("") if err != nil { - slog.Error(fmt.Sprintf("failed to load default template config: %w", err)) + slog.Error("failed to load default template config", "error", err) os.Exit(1) } @@ -46,11 +46,11 @@ func NewLocalWorkerPool() *WorkerPool { defaultTemplateConfig.Import_cache_tree = "" if err := common.SaveConfig(defaultTemplateConfig, templatePath); err != nil { - slog.Error(fmt.Sprintf("failed to save template.json: %w", err)) + slog.Error("failed to save template.json", "error", err) os.Exit(1) } } else { - slog.Error(fmt.Sprintf("failed to stat template path: %w", err)) + slog.Error("failed to stat template path", "error", err) os.Exit(1) } } @@ -99,10 +99,10 @@ func (p *LocalWorkerPoolPlatform) CreateInstance(worker *Worker) error { slog.Error("Failed to get worker config", "workerId", worker.workerId, "error", err) return err } - + // Set worker-specific port cfg.Worker_port = workerPort - + // Save to worker directory configPath := filepath.Join(workerPath, "config.json") if err := common.SaveConfig(cfg, configPath); err != nil { diff --git a/go/boss/lambdastore/store.go b/go/boss/lambdastore/store.go index a43f04e95..8c97564fe 100644 --- a/go/boss/lambdastore/store.go +++ b/go/boss/lambdastore/store.go @@ -262,16 +262,19 @@ func (s *LambdaStore) addToRegistry(funcName string, body io.Reader) error { if err != nil { return fmt.Errorf("failed to create blob writer: %w", err) } - defer func() { - if err := writer.Close(); err != nil { - slog.Error(fmt.Sprintf("warning: failed to close blob writer: %v", err)) - } - }() if _, err := io.Copy(writer, tempFile); err != nil { + // Close writer to release resources (ignore close error since we already have an error) + writer.Close() return fmt.Errorf("failed to upload to blob storage: %w", err) } + // Close the writer to finalize the upload - this is where the actual commit happens + // for many blob storage implementations, so we must check the error + if err := writer.Close(); err != nil { + return fmt.Errorf("failed to finalize blob upload: %w", err) + } + lambdaEntry.Config = cfg if s.eventManager != nil { diff --git a/go/common/config.go b/go/common/config.go index fa90d7188..4f3202183 100644 --- a/go/common/config.go +++ b/go/common/config.go @@ -68,6 +68,20 @@ type Config struct { Features FeaturesConfig `json:"features"` Trace TraceConfig `json:"trace"` Storage StorageConfig `json:"storage"` + Kafka KafkaConfig `json:"kafka"` +} + +type KafkaConfig struct { + // whether to enable the LRU message cache for seek-based replay + Cache_enabled bool `json:"cache_enabled"` + // maximum number of records held in the LRU cache + Cache_size int `json:"cache_size"` + // Kafka consumer session timeout in seconds + Session_timeout_sec int `json:"session_timeout_sec"` + // Kafka consumer heartbeat interval in seconds + Heartbeat_interval_sec int `json:"heartbeat_interval_sec"` + // poll timeout in seconds for each PollFetches call + Poll_timeout_sec int `json:"poll_timeout_sec"` } type DockerConfig struct { @@ -316,6 +330,13 @@ func getDefaultConfigForPatching(olPath string) (*Config, error) { Scratch: "", Code: "", }, + Kafka: KafkaConfig{ + Cache_enabled: true, + Cache_size: 1024, + Session_timeout_sec: 10, + Heartbeat_interval_sec: 3, + Poll_timeout_sec: 1, + }, } return cfg, nil @@ -394,6 +415,8 @@ func checkConf(cfg *Config) error { if cfg.Features.Import_cache != "" { return fmt.Errorf("features.import_cache must be disabled for docker Sandbox") } + } else if cfg.Sandbox == "mock" { + // mock sandbox: no additional requirements } else { return fmt.Errorf("Unknown Sandbox type '%s'", cfg.Sandbox) } @@ -492,3 +515,8 @@ func GetOlPath(ctx *cli.Context) (string, error) { } return filepath.Abs(olPath) } + +// CgroupPoolPath returns the cgroup pool root path for the given OL directory. +func CgroupPoolPath(olPath string) string { + return filepath.Join("/sys/fs/cgroup", filepath.Base(olPath)+"-sandboxes") +} diff --git a/go/common/lambdaConfig.go b/go/common/lambdaConfig.go index 59c0138d1..4ce23bd93 100644 --- a/go/common/lambdaConfig.go +++ b/go/common/lambdaConfig.go @@ -10,6 +10,7 @@ import ( "os" "path/filepath" "regexp" + "strings" "gopkg.in/yaml.v3" ) @@ -42,7 +43,9 @@ type KafkaTrigger struct { // LambdaConfig defines the overall configuration for the lambda function. type LambdaConfig struct { - Triggers Triggers `yaml:"triggers"` // List of HTTP triggers + Triggers Triggers `yaml:"triggers"` // List of HTTP triggers + Environment map[string]string `yaml:"environment"` // Environment variables for the lambda + ReuseSandbox bool `yaml:"reuse-sandbox"` // if true, sandbox is reused across invocations // Additional configurations can be added here. } @@ -54,6 +57,8 @@ func LoadDefaultLambdaConfig() *LambdaConfig { {Method: "*"}, // Default to allow all methods }, }, + Environment: make(map[string]string), + ReuseSandbox: true, } } @@ -87,6 +92,19 @@ func checkLambdaConfig(config *LambdaConfig) error { } } + // Validate environment variables + for key, value := range config.Environment { + if key == "" { + return fmt.Errorf("Environment variable key cannot be empty") + } + // Optionally validate that keys are valid environment variable names + if strings.Contains(key, "=") { + return fmt.Errorf("Environment variable key '%s' cannot contain '='", key) + } + // Value can be empty (that's valid) + _ = value + } + return nil } @@ -104,15 +122,15 @@ func LoadLambdaConfig(codeDir string) (*LambdaConfig, error) { } defer file.Close() - var config LambdaConfig + config := LoadDefaultLambdaConfig() decoder := yaml.NewDecoder(file) - err = decoder.Decode(&config) // Use LambdaConf instead of Conf + err = decoder.Decode(config) // Use LambdaConf instead of Conf if err != nil { return nil, fmt.Errorf("failed to parse YAML file: %v", err) } - return &config, checkLambdaConfig(&config) + return config, checkLambdaConfig(config) } func ExtractConfigFromTarGz(tarPath string) (*LambdaConfig, error) { @@ -143,12 +161,12 @@ func ExtractConfigFromTarGz(tarPath string) (*LambdaConfig, error) { // and ./ matches (./ol.yaml) as tar can encode ./ into filenames under // certain conditions. if filepath.Clean(header.Name) == LambdaConfigFilename { - var config LambdaConfig + config := LoadDefaultLambdaConfig() decoder := yaml.NewDecoder(tr) - if err := decoder.Decode(&config); err != nil { + if err := decoder.Decode(config); err != nil { return nil, fmt.Errorf("failed to parse %s: %w", LambdaConfigFilename, err) } - return &config, checkLambdaConfig(&config) + return config, checkLambdaConfig(config) } } diff --git a/go/common/lambdaConfig_test.go b/go/common/lambdaConfig_test.go index ea776c370..843d3b30b 100644 --- a/go/common/lambdaConfig_test.go +++ b/go/common/lambdaConfig_test.go @@ -9,6 +9,55 @@ import ( "testing" ) +// TestReuseSandbox verifies that the reuse-sandbox field defaults to true +// when not specified, and correctly reflects the value when explicitly set. +func TestReuseSandbox(t *testing.T) { + tests := []struct { + name string + yaml string + expected bool + }{ + { + name: "no ol.yaml — defaults to true", + yaml: "", + expected: true, + }, + { + name: "ol.yaml present but reuse-sandbox is not specified - defaults to true", + yaml: "triggers:\n http:\n - method: \"GET\"\n", + expected: true, + }, + { + name: "reuse-sandbox explicitly set to false", + yaml: "reuse-sandbox: false\n", + expected: false, + }, + { + name: "reuse-sandbox explicitly set to true", + yaml: "reuse-sandbox: true\n", + expected: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dir := t.TempDir() + if tt.yaml != "" { + if err := os.WriteFile(filepath.Join(dir, "ol.yaml"), []byte(tt.yaml), 0644); err != nil { + t.Fatal(err) + } + } + config, err := LoadLambdaConfig(dir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if config.ReuseSandbox != tt.expected { + t.Errorf("expected ReuseSandbox=%v, got %v", tt.expected, config.ReuseSandbox) + } + }) + } +} + // createTestTarGz creates a tar.gz file in memory with the given files func createTestTarGz(t *testing.T, files map[string]string) []byte { var buf bytes.Buffer diff --git a/go/go.mod b/go/go.mod index 21a70ee32..58e168923 100644 --- a/go/go.mod +++ b/go/go.mod @@ -1,6 +1,6 @@ module github.com/open-lambda/open-lambda/go -go 1.24 +go 1.24.0 toolchain go1.24.4 @@ -111,12 +111,12 @@ require ( go.opentelemetry.io/otel/sdk v1.36.0 // indirect go.opentelemetry.io/otel/sdk/metric v1.36.0 // indirect go.opentelemetry.io/otel/trace v1.36.0 // indirect - golang.org/x/crypto v0.39.0 // indirect - golang.org/x/net v0.41.0 // indirect + golang.org/x/crypto v0.45.0 // indirect + golang.org/x/net v0.47.0 // indirect golang.org/x/oauth2 v0.30.0 // indirect - golang.org/x/sync v0.15.0 // indirect - golang.org/x/sys v0.33.0 // indirect - golang.org/x/text v0.26.0 // indirect + golang.org/x/sync v0.18.0 // indirect + golang.org/x/sys v0.38.0 // indirect + golang.org/x/text v0.31.0 // indirect golang.org/x/time v0.11.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect google.golang.org/api v0.235.0 // indirect diff --git a/go/go.sum b/go/go.sum index c080ab5ef..11c94f6c5 100644 --- a/go/go.sum +++ b/go/go.sum @@ -286,8 +286,8 @@ golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPh golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc= golang.org/x/crypto v0.18.0/go.mod h1:R0j02AL6hcrfOiy9T4ZYp/rcWeMxM3L6QYxlOuEG1mg= -golang.org/x/crypto v0.39.0 h1:SHs+kF4LP+f+p14esP5jAoDpHU8Gu/v9lFRK6IT5imM= -golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U= +golang.org/x/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q= +golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= @@ -313,8 +313,8 @@ golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg= golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk= golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY= -golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw= -golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA= +golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY= +golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= @@ -327,8 +327,8 @@ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= -golang.org/x/sync v0.15.0 h1:KWH3jNZsfyT6xfAfKiz6MRNmd46ByHDYaZ7KSkCtdW8= -golang.org/x/sync v0.15.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= +golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I= +golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -344,16 +344,16 @@ golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw= -golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc= +golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k= golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo= golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU= golang.org/x/term v0.16.0/go.mod h1:yn7UURbUtPyrVJPGPq404EukNFxcm/foM+bV/bfcDsY= -golang.org/x/term v0.32.0 h1:DR4lr0TjUs3epypdhTOkMmuF5CDFJ/8pOnbzMZPQ7bg= -golang.org/x/term v0.32.0/go.mod h1:uZG1FhGx848Sqfsq4/DlJr3xGGsYMu/L5GW4abiaEPQ= +golang.org/x/term v0.37.0 h1:8EGAD0qCmHYZg6J17DvsMy9/wJ7/D/4pV/wfnld5lTU= +golang.org/x/term v0.37.0/go.mod h1:5pB4lxRNYYVZuTLmy8oR2BH8dflOR+IbTYFD8fi3254= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= @@ -361,8 +361,8 @@ golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= -golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M= -golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA= +golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM= +golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM= golang.org/x/time v0.11.0 h1:/bpjEDfN9tkoN/ryeYHnv5hcMlc8ncjMcM4XBk5NWV0= golang.org/x/time v0.11.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= diff --git a/go/worker/commands.go b/go/worker/commands.go index b763283bc..2ef49aad9 100644 --- a/go/worker/commands.go +++ b/go/worker/commands.go @@ -16,6 +16,7 @@ import ( "github.com/open-lambda/open-lambda/go/common" "github.com/open-lambda/open-lambda/go/worker/event" + "github.com/open-lambda/open-lambda/go/worker/sandbox/cgroups" "github.com/urfave/cli/v2" ) @@ -39,18 +40,27 @@ func udsGet(requestPath string) (*http.Response, error) { // initCmd corresponds to the "init" command of the admin tool. func initCmd(ctx *cli.Context) error { + if os.Getuid() != 0 { + return fmt.Errorf("'ol worker init' must be run with sudo") + } + olPath, err := common.GetOlPath(ctx) if err != nil { - return err + return fmt.Errorf("init failed to get OL path: %w", err) } if err := common.LoadDefaults(olPath); err != nil { - return err + return fmt.Errorf("init failed to load config defaults: %w", err) } if err := initOLDir(olPath, ctx.String("image"), ctx.Bool("newbase")); err != nil { - return err + return fmt.Errorf("init failed to create OL directory: %w", err) + } + + if err := cgroups.InitPoolRoot(common.CgroupPoolPath(olPath)); err != nil { + return fmt.Errorf("init failed to create cgroup pool: %w", err) } + fmt.Printf("\nYou may optionally modify the defaults here: %s\n\n", filepath.Join(olPath, "config.json")) fmt.Printf("Next start a worker using the \"ol worker up\" command.\n") diff --git a/go/worker/embedded/packagePullerInstaller.py b/go/worker/embedded/packagePullerInstaller.py index cead4c0c3..cc1307bac 100644 --- a/go/worker/embedded/packagePullerInstaller.py +++ b/go/worker/embedded/packagePullerInstaller.py @@ -69,10 +69,11 @@ def f(event): if not alreadyInstalled: try: subprocess.check_output( - ['pip3', 'install', '--no-deps', pkg, '--cache-dir', '/tmp/.cache', '-t', '/host/files']) + ['pip3', 'install', '--no-deps', pkg, '--cache-dir', '/tmp/.cache', '-t', '/host/files'], + stderr=subprocess.STDOUT) except subprocess.CalledProcessError as e: - print(f'pip install failed with error code {e.returncode}') - print(f'Output: {e.output}') + output = e.output.decode('utf-8') if e.output else '' + raise Exception(f'pip install failed for {pkg} (exit code {e.returncode}): {output}') from None name = pkg.split("==")[0] d = deps("/host/files") diff --git a/go/worker/event/cachedKafkaClient.go b/go/worker/event/cachedKafkaClient.go new file mode 100644 index 000000000..b091039aa --- /dev/null +++ b/go/worker/event/cachedKafkaClient.go @@ -0,0 +1,126 @@ +package event + +import ( + "context" + "log/slog" + + "github.com/twmb/franz-go/pkg/kgo" +) + +// cacheKey uniquely identifies a Kafka record by its topic, partition, and offset. +type cacheKey struct { + topic string + partition int32 + offset int64 +} + +// seekState tracks the current seek position when replaying from cache. +type seekState struct { + topic string + partition int32 + offset int64 // next offset to serve from cache +} + +// seekRequest is returned by processMessage when the lambda requests a seek. +type seekRequest struct { + offset int64 +} + +// cachedKafkaClient wraps a KafkaClient and caches records in an LRU map keyed +// by {topic, partition, offset}. When a seek is active, PollFetches serves +// records from the cache. On cache miss, it calls Seek on the underlying +// client so the next poll fetches from the right position. +type cachedKafkaClient struct { + underlying KafkaClient + cache map[cacheKey]*kgo.Record + evictOrder []cacheKey // front = least recently used + maxSize int + seekTarget *seekState +} + +func newCachedKafkaClient(underlying KafkaClient, maxSize int) *cachedKafkaClient { + return &cachedKafkaClient{ + underlying: underlying, + cache: make(map[cacheKey]*kgo.Record), + maxSize: maxSize, + } +} + +// Seek sets the seek target so that subsequent PollFetches calls serve from cache. +func (c *cachedKafkaClient) Seek(topic string, partition int32, offset int64) { + c.seekTarget = &seekState{topic: topic, partition: partition, offset: offset} +} + +// PollFetches serves from cache when seeking, otherwise delegates to the underlying client. +func (c *cachedKafkaClient) PollFetches(ctx context.Context) kgo.Fetches { + if c.seekTarget != nil { + key := cacheKey{ + topic: c.seekTarget.topic, + partition: c.seekTarget.partition, + offset: c.seekTarget.offset, + } + record, ok := c.cache[key] + if ok { + c.touchLRU(key) + c.seekTarget.offset++ + return makeSingleRecordFetches(record) + } + // Cache miss — tell the underlying client to fetch from this offset. + // The next normal PollFetches will get records starting here. + slog.Info("Seek cache miss, setting offset on underlying client", + "topic", c.seekTarget.topic, + "partition", c.seekTarget.partition, + "offset", c.seekTarget.offset) + c.underlying.Seek(c.seekTarget.topic, c.seekTarget.partition, c.seekTarget.offset) + c.seekTarget = nil + } + + fetches := c.underlying.PollFetches(ctx) + fetches.EachRecord(func(record *kgo.Record) { + c.put(cacheKey{topic: record.Topic, partition: record.Partition, offset: record.Offset}, record) + }) + return fetches +} + +func (c *cachedKafkaClient) Close() { + c.underlying.Close() +} + +// put adds a record to the cache, evicting the LRU entry if at capacity. +func (c *cachedKafkaClient) put(key cacheKey, record *kgo.Record) { + if _, exists := c.cache[key]; exists { + c.touchLRU(key) + return + } + if len(c.cache) >= c.maxSize { + evictKey := c.evictOrder[0] + c.evictOrder = c.evictOrder[1:] + delete(c.cache, evictKey) + } + c.cache[key] = record + c.evictOrder = append(c.evictOrder, key) +} + +// touchLRU moves a key to the back of the eviction order (most recently used). +func (c *cachedKafkaClient) touchLRU(key cacheKey) { + for i, k := range c.evictOrder { + if k == key { + c.evictOrder = append(c.evictOrder[:i], c.evictOrder[i+1:]...) + c.evictOrder = append(c.evictOrder, key) + return + } + } +} + +// makeSingleRecordFetches wraps a single record into the kgo.Fetches structure. +func makeSingleRecordFetches(record *kgo.Record) kgo.Fetches { + return kgo.Fetches{{ + Topics: []kgo.FetchTopic{{ + Topic: record.Topic, + Partitions: []kgo.FetchPartition{{ + Partition: record.Partition, + Records: []*kgo.Record{record}, + }}, + }}, + }} +} diff --git a/go/worker/event/kafkaServer.go b/go/worker/event/kafkaServer.go index 993dffc55..7b590f984 100644 --- a/go/worker/event/kafkaServer.go +++ b/go/worker/event/kafkaServer.go @@ -9,6 +9,7 @@ import ( "log/slog" "net/http" "net/http/httptest" + "strconv" "strings" "sync" "time" @@ -21,24 +22,60 @@ import ( type KafkaClient interface { PollFetches(context.Context) kgo.Fetches + Seek(topic string, partition int32, offset int64) Close() } +// kgoClientWrapper wraps *kgo.Client to implement KafkaClient. +type kgoClientWrapper struct { + client *kgo.Client +} + +func (w *kgoClientWrapper) PollFetches(ctx context.Context) kgo.Fetches { + return w.client.PollFetches(ctx) +} + +func (w *kgoClientWrapper) Seek(topic string, partition int32, offset int64) { + w.client.SetOffsets(map[string]map[int32]kgo.EpochOffset{ + topic: {partition: {Offset: offset}}, + }) +} + +func (w *kgoClientWrapper) Close() { + w.client.Close() +} + +// LambdaInvoker abstracts the lambda invocation layer for testability +type LambdaInvoker interface { + Invoke(lambdaName string, w http.ResponseWriter, r *http.Request) +} + +// lambdaMgrInvoker wraps *lambda.LambdaMgr to implement LambdaInvoker +type lambdaMgrInvoker struct { + mgr *lambda.LambdaMgr +} + +func (i *lambdaMgrInvoker) Invoke(lambdaName string, w http.ResponseWriter, r *http.Request) { + f := i.mgr.Get(lambdaName) + f.Invoke(w, r) +} + // LambdaKafkaConsumer manages Kafka consumption for a specific lambda function type LambdaKafkaConsumer struct { - consumerName string // Unique name for this consumer - lambdaName string // lambda function name - kafkaTrigger *common.KafkaTrigger - client KafkaClient // kgo.client implements the KafkaClient interface - lambdaManager *lambda.LambdaMgr // Reference to lambda manager for direct calls - stopChan chan struct{} // Shutdown signal for this consumer + consumerName string // Unique name for this consumer + lambdaName string // lambda function name + kafkaTrigger *common.KafkaTrigger + client KafkaClient // used for PollFetches/Close/Seek + invoker LambdaInvoker // Abstraction for lambda invocation + stopChan chan struct{} // Shutdown signal for this consumer // When this channel is closed, the goroutine for the consumer exits + errorCount int // Number of non-timeout Kafka client errors encountered } // KafkaManager manages multiple lambda-specific Kafka consumers type KafkaManager struct { lambdaConsumers map[string]*LambdaKafkaConsumer // lambdaName -> consumer - lambdaManager *lambda.LambdaMgr // Reference to lambda manager + invoker LambdaInvoker // Abstraction for lambda invocation mu sync.Mutex // Protects lambdaConsumers map } @@ -52,13 +89,15 @@ func (km *KafkaManager) newLambdaKafkaConsumer(consumerName string, lambdaName s return nil, fmt.Errorf("no topics configured for lambda %s", lambdaName) } + kafkaCfg := common.Conf.Kafka + // Setup kgo client options opts := []kgo.Opt{ kgo.SeedBrokers(trigger.BootstrapServers...), kgo.ConsumerGroup(trigger.GroupId), kgo.ConsumeTopics(trigger.Topics...), - kgo.SessionTimeout(10 * time.Second), - kgo.HeartbeatInterval(3 * time.Second), + kgo.SessionTimeout(time.Duration(kafkaCfg.Session_timeout_sec) * time.Second), + kgo.HeartbeatInterval(time.Duration(kafkaCfg.Heartbeat_interval_sec) * time.Second), } // Use trigger-specific offset reset or default to latest @@ -74,13 +113,17 @@ func (km *KafkaManager) newLambdaKafkaConsumer(consumerName string, lambdaName s return nil, fmt.Errorf("failed to create Kafka client for lambda %s: %w", lambdaName, err) } + var kafkaClient KafkaClient = &kgoClientWrapper{client: client} + if kafkaCfg.Cache_enabled { + kafkaClient = newCachedKafkaClient(kafkaClient, kafkaCfg.Cache_size) + } return &LambdaKafkaConsumer{ - consumerName: consumerName, - lambdaName: lambdaName, - kafkaTrigger: trigger, - client: client, - lambdaManager: km.lambdaManager, - stopChan: make(chan struct{}), + consumerName: consumerName, + lambdaName: lambdaName, + kafkaTrigger: trigger, + client: kafkaClient, + invoker: km.invoker, + stopChan: make(chan struct{}), }, nil } @@ -88,7 +131,7 @@ func (km *KafkaManager) newLambdaKafkaConsumer(consumerName string, lambdaName s func NewKafkaManager(lambdaManager *lambda.LambdaMgr) (*KafkaManager, error) { manager := &KafkaManager{ lambdaConsumers: make(map[string]*LambdaKafkaConsumer), - lambdaManager: lambdaManager, + invoker: &lambdaMgrInvoker{mgr: lambdaManager}, } slog.Info("Kafka manager initialized") @@ -116,7 +159,7 @@ func (lkc *LambdaKafkaConsumer) consumeLoop() { slog.Info("Stopping Kafka consumer for lambda", "lambda", lkc.lambdaName) return default: - ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), time.Duration(common.Conf.Kafka.Poll_timeout_sec)*time.Second) fetches := lkc.client.PollFetches(ctx) cancel() @@ -129,6 +172,7 @@ func (lkc *LambdaKafkaConsumer) consumeLoop() { continue } + lkc.errorCount++ // TODO: Surface Kafka consumer errors to lambda developers by invoking an error // handler lambda function. Could allow lambdas to specify an onError callback in // ol.yaml that gets invoked with error details. @@ -139,8 +183,9 @@ func (lkc *LambdaKafkaConsumer) consumeLoop() { continue } - // Process each record - fetches.EachRecord(func(record *kgo.Record) { + // Process each record. Manual iteration (instead of EachRecord) lets + // us break out mid-batch when a seek is requested. + for _, record := range fetches.Records() { slog.Info("Received Kafka message for lambda", "consumer", lkc.consumerName, "lambda", lkc.lambdaName, @@ -148,26 +193,34 @@ func (lkc *LambdaKafkaConsumer) consumeLoop() { "partition", record.Partition, "offset", record.Offset, "size", len(record.Value)) - lkc.processMessage(record) - }) + if seek := lkc.processMessage(record); seek != nil { + lkc.client.Seek(record.Topic, record.Partition, seek.offset) + break // next PollFetches will serve from cache + } + } } } } -// processMessage handles a single Kafka message by invoking the lambda function directly -func (lkc *LambdaKafkaConsumer) processMessage(record *kgo.Record) { +// processMessage handles a single Kafka message by invoking the lambda function directly. +// If the lambda returns an X-Kafka-Seek-Offset header, the corresponding seekRequest is returned. +func (lkc *LambdaKafkaConsumer) processMessage(record *kgo.Record) *seekRequest { t := common.T0("kafka-message-processing") defer t.T1() // Create synthetic HTTP request from Kafka message - req, err := http.NewRequest("POST", "/", bytes.NewReader(record.Value)) + // Path must be /run// for the Python runtime to parse correctly + requestPath := fmt.Sprintf("/run/%s/", lkc.lambdaName) + req, err := http.NewRequest("POST", requestPath, bytes.NewReader(record.Value)) if err != nil { slog.Error("Failed to create request for lambda invocation", "lambda", lkc.lambdaName, "error", err, "topic", record.Topic) - return + return nil } + // RequestURI must be set explicitly for synthetic requests (http.NewRequest doesn't set it) + req.RequestURI = requestPath // Set headers with Kafka metadata (The X- prefix indicates a custom non-standard header) req.Header.Set("Content-Type", "application/json") @@ -177,13 +230,10 @@ func (lkc *LambdaKafkaConsumer) processMessage(record *kgo.Record) { req.Header.Set("X-Kafka-Group-Id", lkc.kafkaTrigger.GroupId) // Create response recorder to capture lambda output. - // TODO: Capture and log the lambda response body using httptest's response recorder - // for kafka triggered lambda invocations. w := httptest.NewRecorder() - // Get lambda function and invoke directly - lambdaFunc := lkc.lambdaManager.Get(lkc.lambdaName) - lambdaFunc.Invoke(w, req) + // Invoke the lambda function directly + lkc.invoker.Invoke(lkc.lambdaName, w, req) // Log the result slog.Info("Kafka message processed via direct invocation", @@ -193,6 +243,26 @@ func (lkc *LambdaKafkaConsumer) processMessage(record *kgo.Record) { "partition", record.Partition, "offset", record.Offset, "status", w.Code) + + // Check if the lambda requested a seek via response header + if seekStr := w.Header().Get("X-Kafka-Seek-Offset"); seekStr != "" { + seekOffset, err := strconv.ParseInt(seekStr, 10, 64) + if err != nil { + slog.Warn("Invalid X-Kafka-Seek-Offset header", + "lambda", lkc.lambdaName, + "value", seekStr, + "error", err) + return nil + } + slog.Info("Lambda requested seek", + "lambda", lkc.lambdaName, + "topic", record.Topic, + "partition", record.Partition, + "current_offset", record.Offset, + "seek_offset", seekOffset) + return &seekRequest{offset: seekOffset} + } + return nil } // cleanup closes the kgo client diff --git a/go/worker/event/kafkaServer_test.go b/go/worker/event/kafkaServer_test.go new file mode 100644 index 000000000..4311e1514 --- /dev/null +++ b/go/worker/event/kafkaServer_test.go @@ -0,0 +1,470 @@ +package event + +import ( + "context" + "fmt" + "io" + "net/http" + "os" + "reflect" + "sync" + "sync/atomic" + "testing" + + "github.com/open-lambda/open-lambda/go/common" + "github.com/twmb/franz-go/pkg/kgo" +) + +func TestMain(m *testing.M) { + // Initialize common.Conf so that common.T0/T1 (latency tracking) doesn't panic + common.Conf = &common.Config{ + Kafka: common.KafkaConfig{ + Cache_enabled: true, + Cache_size: 1024, + Session_timeout_sec: 10, + Heartbeat_interval_sec: 3, + Poll_timeout_sec: 1, + }, + } + os.Exit(m.Run()) +} + +// --- Mocks --- + +// MockKafkaClient implements KafkaClient for testing. +// +// Instead of exposing pollFetchesFunc directly, tests enqueue responses via +// Send and SendError. The mock serves them in FIFO order and returns empty +// fetches once the queue is drained. This keeps polling/sequencing logic out +// of individual tests. +// +// The Drained channel (when set) is closed the first time PollFetches is called +// after the queue is empty. Because the consume loop calls PollFetches only +// after finishing the previous iteration's processing, a receive on Drained +// guarantees all enqueued records have been fully processed. +type MockKafkaClient struct { + mu sync.Mutex + queue []kgo.Fetches + callCount int + closeCalled atomic.Bool + Drained chan struct{} // closed when all queued fetches have been consumed and processed + drainedSignaled bool +} + +// Send enqueues records that will be returned by the next PollFetches call. +func (m *MockKafkaClient) Send(records ...*kgo.Record) { + m.mu.Lock() + defer m.mu.Unlock() + m.queue = append(m.queue, makeFetches(records...)) +} + +// SendError enqueues a fetch error that will be returned by the next PollFetches call. +func (m *MockKafkaClient) SendError(topic string, partition int32, err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.queue = append(m.queue, makeErrorFetches(topic, partition, err)) +} + +// PollFetches returns the next queued fetch, or empty fetches if the queue is +// drained. When the queue is empty and Drained is set, it closes Drained to +// signal that all prior records have been processed. +func (m *MockKafkaClient) PollFetches(ctx context.Context) kgo.Fetches { + m.mu.Lock() + defer m.mu.Unlock() + if m.callCount < len(m.queue) { + f := m.queue[m.callCount] + m.callCount++ + return f + } + if !m.drainedSignaled && m.Drained != nil { + close(m.Drained) + m.drainedSignaled = true + } + return kgo.Fetches{} +} + +func (m *MockKafkaClient) Seek(topic string, partition int32, offset int64) { + // No-op for mock — tests control what PollFetches returns via Send/SendError +} + +func (m *MockKafkaClient) Close() { + m.closeCalled.Store(true) +} + +// MockLambdaInvoker implements LambdaInvoker for testing. +// When respondFunc is set, it is called with the ResponseWriter and invocation +// index instead of the default w.WriteHeader(200). This lets tests simulate +// custom response headers (e.g., X-Kafka-Seek-Offset) on specific invocations. +type MockLambdaInvoker struct { + mu sync.Mutex + invocations []invokeRecord + respondFunc func(w http.ResponseWriter, invocationIndex int) +} + +// invokeRecord captures the relevant fields from a lambda invocation in simple, +// comparable types. Tests can build an expected invokeRecord and compare it +// directly with reflect.DeepEqual instead of asserting each field individually. +type invokeRecord struct { + LambdaName string + Method string + Path string + RequestURI string + Body string + Headers map[string]string +} + +func (m *MockLambdaInvoker) Invoke(lambdaName string, w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + + // Flatten headers into a simple map for easy comparison in assertions + headers := map[string]string{} + for key := range r.Header { + headers[key] = r.Header.Get(key) + } + + m.mu.Lock() + idx := len(m.invocations) + m.invocations = append(m.invocations, invokeRecord{ + LambdaName: lambdaName, + Method: r.Method, + Path: r.URL.Path, + RequestURI: r.RequestURI, + Body: string(body), + Headers: headers, + }) + respondFunc := m.respondFunc + m.mu.Unlock() + + if respondFunc != nil { + respondFunc(w, idx) + } else { + w.WriteHeader(http.StatusOK) + } +} + +func (m *MockLambdaInvoker) getInvocations() []invokeRecord { + m.mu.Lock() + defer m.mu.Unlock() + cp := make([]invokeRecord, len(m.invocations)) + copy(cp, m.invocations) + return cp +} + +// --- Helpers --- + +// makeFetches converts flat kgo.Records into the nested kgo.Fetches structure. +// +// franz-go's PollFetches returns a deeply nested type that mirrors how Kafka +// brokers organize data: +// +// Fetches -> []Fetch -> []FetchTopic -> []FetchPartition -> []*Record +// +// Records are grouped by topic and then by partition. This helper handles that +// grouping automatically so tests can think in terms of simple records rather +// than the broker-level wire format. +func makeFetches(records ...*kgo.Record) kgo.Fetches { + if len(records) == 0 { + return kgo.Fetches{} + } + // Group records by topic+partition + type key struct { + topic string + partition int32 + } + groups := map[key][]*kgo.Record{} + for _, r := range records { + k := key{r.Topic, r.Partition} + groups[k] = append(groups[k], r) + } + + topicMap := map[string][]kgo.FetchPartition{} + for k, recs := range groups { + topicMap[k.topic] = append(topicMap[k.topic], kgo.FetchPartition{ + Partition: k.partition, + Records: recs, + }) + } + + var topics []kgo.FetchTopic + for topic, partitions := range topicMap { + topics = append(topics, kgo.FetchTopic{ + Topic: topic, + Partitions: partitions, + }) + } + return kgo.Fetches{{Topics: topics}} +} + +func makeErrorFetches(topic string, partition int32, err error) kgo.Fetches { + return kgo.Fetches{{ + Topics: []kgo.FetchTopic{{ + Topic: topic, + Partitions: []kgo.FetchPartition{{ + Partition: partition, + Err: err, + }}, + }}, + }} +} + +// setupConsumerHarness creates the full test harness for exercising the consumer's +// consumeLoop. It mocks both sides of the consumer: +// +// - Above the consumer (Kafka broker layer): MockKafkaClient replaces the real +// Kafka connection so tests can enqueue records and errors without a broker. +// - Below the consumer (lambda invocation layer): MockLambdaInvoker replaces the +// real lambda invocation path so tests can capture and assert on HTTP requests. +// +// The consumer itself is real — it runs the actual consumeLoop logic, so tests +// exercise the full record-processing and error-handling pipeline. +func setupConsumerHarness(lambdaName string) (*MockKafkaClient, *MockLambdaInvoker, *LambdaKafkaConsumer) { + // Mock above: fake Kafka broker + client := &MockKafkaClient{Drained: make(chan struct{})} + // Mock below: fake lambda invocation + invoker := &MockLambdaInvoker{} + + consumer := &LambdaKafkaConsumer{ + consumerName: lambdaName + "-0", + lambdaName: lambdaName, + kafkaTrigger: &common.KafkaTrigger{GroupId: "lambda-" + lambdaName}, + client: client, + invoker: invoker, + stopChan: make(chan struct{}), + } + return client, invoker, consumer +} + +// runConsumeLoop starts consumeLoop in a goroutine and returns a stop function +// that signals shutdown and waits for the goroutine to exit. Callers should +// <-mockClient.Drained before stop() — Drained closes once all enqueued +// records have been fully processed, making it safe to assert on results. +func runConsumeLoop(consumer *LambdaKafkaConsumer) (stop func()) { + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + consumer.consumeLoop() + }() + return func() { + close(consumer.stopChan) + wg.Wait() + } +} + +// --- Tests --- + +func TestConsumeLoop_ProcessesRecords(t *testing.T) { + mockClient, invoker, consumer := setupConsumerHarness("my-lambda") + mockClient.Send(&kgo.Record{ + Topic: "orders", Partition: 3, Offset: 99, + Value: []byte(`{"orderId": 42}`), + }) + stop := runConsumeLoop(consumer) + <-mockClient.Drained + stop() + + invocations := invoker.getInvocations() + if len(invocations) != 1 { + t.Fatalf("Expected 1 invocation, got %d", len(invocations)) + } + + expected := invokeRecord{ + LambdaName: "my-lambda", + Method: "POST", + Path: "/run/my-lambda/", + RequestURI: "/run/my-lambda/", + Body: `{"orderId": 42}`, + Headers: map[string]string{ + "Content-Type": "application/json", + "X-Kafka-Topic": "orders", + "X-Kafka-Partition": "3", + "X-Kafka-Offset": "99", + "X-Kafka-Group-Id": "lambda-my-lambda", + }, + } + if !reflect.DeepEqual(invocations[0], expected) { + t.Errorf("Invocation mismatch:\n got: %+v\n want: %+v", invocations[0], expected) + } +} + +func TestConsumeLoop_ContinuesThroughErrors(t *testing.T) { + mockClient, invoker, consumer := setupConsumerHarness("test-lambda") + // Poll sequence: deadline-exceeded errors (silently skipped), then a real + // error (counted), then a valid record. The loop should survive all of them. + mockClient.SendError("topic", 0, context.DeadlineExceeded) + mockClient.SendError("topic", 0, context.DeadlineExceeded) + mockClient.SendError("topic", 0, fmt.Errorf("broker unreachable")) + mockClient.Send(&kgo.Record{ + Topic: "topic", Partition: 0, Offset: 1, Value: []byte("survived"), + }) + stop := runConsumeLoop(consumer) + <-mockClient.Drained + stop() + + invocations := invoker.getInvocations() + if len(invocations) != 1 { + t.Fatalf("Expected 1 invocation after errors, got %d", len(invocations)) + } + + expected := invokeRecord{ + LambdaName: "test-lambda", + Method: "POST", + Path: "/run/test-lambda/", + RequestURI: "/run/test-lambda/", + Body: "survived", + Headers: map[string]string{ + "Content-Type": "application/json", + "X-Kafka-Topic": "topic", + "X-Kafka-Partition": "0", + "X-Kafka-Offset": "1", + "X-Kafka-Group-Id": "lambda-test-lambda", + }, + } + if !reflect.DeepEqual(invocations[0], expected) { + t.Errorf("Invocation mismatch:\n got: %+v\n want: %+v", invocations[0], expected) + } + + // Only real errors should be counted; DeadlineExceeded should be ignored + if consumer.errorCount != 1 { + t.Errorf("Expected 1 error counted, got %d", consumer.errorCount) + } +} + +func TestUnregister(t *testing.T) { + manager := &KafkaManager{ + lambdaConsumers: make(map[string]*LambdaKafkaConsumer), + } + + mockClient := &MockKafkaClient{} + manager.lambdaConsumers["test-lambda-0"] = &LambdaKafkaConsumer{ + consumerName: "test-lambda-0", + lambdaName: "test-lambda", + client: mockClient, + stopChan: make(chan struct{}), + } + + manager.UnregisterLambdaKafkaTriggers("test-lambda") + + if len(manager.lambdaConsumers) != 0 { + t.Errorf("Expected 0 consumers, got %d", len(manager.lambdaConsumers)) + } + if !mockClient.closeCalled.Load() { + t.Error("Expected Close to be called on client") + } +} + +// --- cachedKafkaClient unit tests --- + +func TestCachedClient_CachesRecords(t *testing.T) { + mock := &MockKafkaClient{Drained: make(chan struct{})} + mock.Send( + &kgo.Record{Topic: "t", Partition: 0, Offset: 0, Value: []byte("a")}, + &kgo.Record{Topic: "t", Partition: 0, Offset: 1, Value: []byte("b")}, + ) + + cached := newCachedKafkaClient(mock, 100) + cached.PollFetches(context.Background()) + + // Both records should now be in the cache + if _, ok := cached.cache[cacheKey{"t", 0, 0}]; !ok { + t.Error("Expected offset 0 to be cached") + } + if _, ok := cached.cache[cacheKey{"t", 0, 1}]; !ok { + t.Error("Expected offset 1 to be cached") + } +} + +func TestCachedClient_SeekCacheHit(t *testing.T) { + mock := &MockKafkaClient{Drained: make(chan struct{})} + mock.Send( + &kgo.Record{Topic: "t", Partition: 0, Offset: 10, Value: []byte("ten")}, + &kgo.Record{Topic: "t", Partition: 0, Offset: 11, Value: []byte("eleven")}, + &kgo.Record{Topic: "t", Partition: 0, Offset: 12, Value: []byte("twelve")}, + ) + + cached := newCachedKafkaClient(mock, 100) + // Populate the cache + cached.PollFetches(context.Background()) + + // Seek to offset 10 + cached.Seek("t", 0, 10) + + // Each PollFetches should return the next cached record + f1 := cached.PollFetches(context.Background()) + records1 := f1.Records() + if len(records1) != 1 || records1[0].Offset != 10 { + t.Fatalf("Expected offset 10, got %v", records1) + } + + f2 := cached.PollFetches(context.Background()) + records2 := f2.Records() + if len(records2) != 1 || records2[0].Offset != 11 { + t.Fatalf("Expected offset 11, got %v", records2) + } + + f3 := cached.PollFetches(context.Background()) + records3 := f3.Records() + if len(records3) != 1 || records3[0].Offset != 12 { + t.Fatalf("Expected offset 12, got %v", records3) + } +} + +func TestCachedClient_SeekCacheMiss(t *testing.T) { + mock := &MockKafkaClient{Drained: make(chan struct{})} + mock.Send( + &kgo.Record{Topic: "t", Partition: 0, Offset: 5, Value: []byte("five")}, + ) + // After SetOffset, the next PollFetches returns from the new position + mock.Send( + &kgo.Record{Topic: "t", Partition: 0, Offset: 99, Value: []byte("ninety-nine")}, + ) + + cached := newCachedKafkaClient(mock, 100) + // Populate cache with offset 5 + cached.PollFetches(context.Background()) + + // Seek to offset 99 which is not in cache + cached.Seek("t", 0, 99) + + // Cache miss clears seek and calls SetOffset on underlying. + // The same PollFetches call falls through to normal polling. + fetches := cached.PollFetches(context.Background()) + records := fetches.Records() + if len(records) != 1 || records[0].Offset != 99 { + t.Fatalf("Expected offset 99 from underlying after cache miss, got %v", records) + } + + // Seek should be cleared after cache miss + if cached.seekTarget != nil { + t.Error("Expected seekTarget to be nil after cache miss") + } + + // The fetched record should now be cached + if _, ok := cached.cache[cacheKey{"t", 0, 99}]; !ok { + t.Error("Expected offset 99 to be cached after fetch") + } +} + +func TestCachedClient_LRUEviction(t *testing.T) { + mock := &MockKafkaClient{Drained: make(chan struct{})} + mock.Send( + &kgo.Record{Topic: "t", Partition: 0, Offset: 0, Value: []byte("a")}, + &kgo.Record{Topic: "t", Partition: 0, Offset: 1, Value: []byte("b")}, + &kgo.Record{Topic: "t", Partition: 0, Offset: 2, Value: []byte("c")}, + ) + + // Cache can only hold 2 records + cached := newCachedKafkaClient(mock, 2) + cached.PollFetches(context.Background()) + + // Offset 0 should have been evicted (LRU), offsets 1 and 2 should remain + if _, ok := cached.cache[cacheKey{"t", 0, 0}]; ok { + t.Error("Expected offset 0 to be evicted") + } + if _, ok := cached.cache[cacheKey{"t", 0, 1}]; !ok { + t.Error("Expected offset 1 to be cached") + } + if _, ok := cached.cache[cacheKey{"t", 0, 2}]; !ok { + t.Error("Expected offset 2 to be cached") + } +} diff --git a/go/worker/event/sockServer.go b/go/worker/event/sockServer.go index 51ceed985..1bb2cdbb0 100644 --- a/go/worker/event/sockServer.go +++ b/go/worker/event/sockServer.go @@ -104,15 +104,16 @@ func (server *SOCKServer) Create(w http.ResponseWriter, _ []string, args map[str } } - if parent != nil && parent.GetRuntimeType() != rtType { + if parent != nil && parent.Meta().Runtime != rtType { return fmt.Errorf("Parent and child have different runtimes") } meta := &sandbox.SandboxMeta{ + Runtime: rtType, Installs: packages, } - c, err := server.sbPool.Create(parent, leaf, codeDir, scratchDir, meta, rtType) + c, err := server.sbPool.Create(parent, leaf, codeDir, scratchDir, meta) if err != nil { return err } diff --git a/go/worker/helpers.go b/go/worker/helpers.go index fdbdd6171..9a882e1bd 100644 --- a/go/worker/helpers.go +++ b/go/worker/helpers.go @@ -21,6 +21,7 @@ import ( "github.com/open-lambda/open-lambda/go/worker/embedded" ) + func initOLBaseDir(baseDir string, dockerBaseImage string) error { if dockerBaseImage == "" { dockerBaseImage = "ol-wasm" @@ -279,8 +280,8 @@ func runningToStoppedClean() error { // It cleans up cgroups and mounts associated with the OpenLambda instance at `olPath`. // Returns errors encountered during cleanup operations. func stoppedDirtyToStoppedClean(olPath string) error { - // Clean up cgroups associated with sandboxes - cgRoot := filepath.Join("/sys", "fs", "cgroup", filepath.Base(olPath)+"-sandboxes") + // Clean up child cgroups, preserving the pool root + cgRoot := common.CgroupPoolPath(olPath) fmt.Printf("Attempting to clean up cgroups at %s\n", cgRoot) cgroupErrorCount := 0 @@ -302,7 +303,6 @@ func stoppedDirtyToStoppedClean(olPath string) error { } kill := filepath.Join(cgRoot, "cgroup.kill") if err := os.WriteFile(kill, []byte(fmt.Sprintf("%d", 1)), os.ModeAppend); err != nil { - // Print an error if killing processes in the cgroup fails. fmt.Printf("Could not kill processes in cgroup: %s\n", err.Error()) cgroupErrorCount += 1 } @@ -311,17 +311,11 @@ func stoppedDirtyToStoppedClean(olPath string) error { cg := filepath.Join(cgRoot, file.Name()) fmt.Printf("Attempting to remove %s\n", cg) if err := syscall.Rmdir(cg); err != nil { - // Print an error if removing a cgroup fails. - fmt.Printf("could not remove cgroup: %s", err.Error()) + fmt.Printf("could not remove cgroup: %s\n", err.Error()) cgroupErrorCount += 1 } } } - if err := syscall.Rmdir(cgRoot); err != nil { - // Print an error if removing the cgroup root directory fails. - fmt.Printf("could not remove cgroup root: %s", err.Error()) - cgroupErrorCount += 1 - } } sandboxErrorCount := 0 diff --git a/go/worker/lambda/handlerPuller.go b/go/worker/lambda/handlerPuller.go index 2c9cb73c6..8c4b75d83 100644 --- a/go/worker/lambda/handlerPuller.go +++ b/go/worker/lambda/handlerPuller.go @@ -7,7 +7,6 @@ import ( "io" "os" "os/exec" - "path/filepath" "strings" "sync" "time" @@ -23,8 +22,6 @@ import ( var errNotFound404 = errors.New("lambda not found in blob store") -var RT_UNKNOWN common.RuntimeType - type HandlerPuller struct { bucket *blob.Bucket dirCache sync.Map // key=lambda name, value=*CacheEntry @@ -34,7 +31,6 @@ type HandlerPuller struct { type CacheEntry struct { version time.Time // blob modification time path string - runtime common.RuntimeType } func NewHandlerPuller(dirMaker *common.DirMaker) (*HandlerPuller, error) { @@ -67,12 +63,12 @@ func NewHandlerPuller(dirMaker *common.DirMaker) (*HandlerPuller, error) { }, nil } -func (cp *HandlerPuller) Pull(name string) (common.RuntimeType, string, error) { +func (cp *HandlerPuller) Pull(name string) (string, error) { t := common.T0("pull-lambda") defer t.T1() if err := common.ValidateFunctionName(name); err != nil { - return RT_UNKNOWN, "", err + return "", err } key := name + common.LambdaFileExtension @@ -81,71 +77,62 @@ func (cp *HandlerPuller) Pull(name string) (common.RuntimeType, string, error) { if err == nil { version := attrs.ModTime if cached := cp.getCache(name); cached != nil && cached.version.Equal(version) { - return cached.runtime, cached.path, nil + return cached.path, nil } } - rt, dir, err := cp.pullFromBlob(key, name) + dir, err := cp.pullFromBlob(key, name) if err == nil { var version time.Time if attrs != nil { version = attrs.ModTime } - cp.putCache(name, version, dir, rt) - return rt, dir, nil + cp.putCache(name, version, dir) + return dir, nil } else if err != errNotFound404 { - return RT_UNKNOWN, "", err + return "", err } - return RT_UNKNOWN, "", fmt.Errorf( + return "", fmt.Errorf( "lambda %q not found in blob store (bucket=%q, key=%q)", name, common.Conf.Registry, key, ) } -func (cp *HandlerPuller) pullFromBlob(key, lambdaName string) (common.RuntimeType, string, error) { +func (cp *HandlerPuller) pullFromBlob(key, lambdaName string) (string, error) { ctx := context.Background() reader, err := cp.bucket.NewReader(ctx, key, nil) if err != nil { if gcerrors.Code(err) == gcerrors.NotFound { - return RT_UNKNOWN, "", errNotFound404 + return "", errNotFound404 } - return RT_UNKNOWN, "", err + return "", err } defer reader.Close() tmpFile, err := os.CreateTemp("", lambdaName+"_blob") - if err != nil { - return RT_UNKNOWN, "", err + return "", err } tmpPath := tmpFile.Name() if _, err := io.Copy(tmpFile, reader); err != nil { tmpFile.Close() - return RT_UNKNOWN, "", err + return "", err } tmpFile.Close() defer os.Remove(tmpPath) targetDir := cp.dirMaker.Get(lambdaName) if err := os.MkdirAll(targetDir, 0755); err != nil { - return RT_UNKNOWN, "", err + return "", err } cmd := exec.Command("tar", "-xzf", tmpPath, "--directory", targetDir) if output, err := cmd.CombinedOutput(); err != nil { - return RT_UNKNOWN, "", fmt.Errorf("tar extract failed: %v :: %s", err, output) + return "", fmt.Errorf("tar extract failed: %v :: %s", err, output) } - var rt common.RuntimeType - if _, err := os.Stat(filepath.Join(targetDir, "f.py")); err == nil { - rt = common.RT_PYTHON - } else if _, err := os.Stat(filepath.Join(targetDir, "f.bin")); err == nil { - rt = common.RT_NATIVE - } else { - return RT_UNKNOWN, "", fmt.Errorf("runtime type not found in extracted archive") - } - return rt, targetDir, nil + return targetDir, nil } func (cp *HandlerPuller) Reset(name string) { @@ -159,10 +146,10 @@ func (cp *HandlerPuller) getCache(name string) *CacheEntry { } return entry.(*CacheEntry) } -func (cp *HandlerPuller) putCache(name string, version time.Time, path string, runtime common.RuntimeType) { +func (cp *HandlerPuller) putCache(name string, version time.Time, path string) { // Clean up old cache entry if it exists if old := cp.getCache(name); old != nil && old.path != path { os.RemoveAll(old.path) } - cp.dirCache.Store(name, &CacheEntry{version, path, runtime}) + cp.dirCache.Store(name, &CacheEntry{version, path}) } diff --git a/go/worker/lambda/lambdaFunction.go b/go/worker/lambda/lambdaFunction.go index f625ad43e..21ff11825 100644 --- a/go/worker/lambda/lambdaFunction.go +++ b/go/worker/lambda/lambdaFunction.go @@ -18,8 +18,11 @@ import ( ) type FunctionMeta struct { - Sandbox *sandbox.SandboxMeta `json:"sandbox"` // Existing sandbox metadata - Config *common.LambdaConfig `json:"config"` // New Lambda config (from YAML) + // user-specified config (via ol.yaml) + Config *common.LambdaConfig `json:"config"` + // container-specific settings, inferred by file contents + // (e.g., do we have Python or native code? what is in requirements.txt?) + Sandbox *sandbox.SandboxMeta `json:"sandbox"` } // LambdaFunc represents a single lambda function (the code) @@ -27,8 +30,6 @@ type LambdaFunc struct { lmgr *LambdaMgr name string - rtType common.RuntimeType - // lambda code lastPull *time.Time codeDir string @@ -72,36 +73,58 @@ func (f *LambdaFunc) printf(format string, args ...any) { slog.Info(fmt.Sprintf("%s [FUNC %s]", strings.TrimRight(msg, "\n"), f.name)) } -// parseMeta reads in a requirements.txt file that was built from pip-compile +// parseMeta constructs a FunctionMeta based on contents of a code +// directory, such as an ol.yaml and requirements.txt (generated by +// pip-compile) func parseMeta(codeDir string) (*FunctionMeta, error) { sandboxMeta := &sandbox.SandboxMeta{ Installs: []string{}, Imports: []string{}, } - path := filepath.Join(codeDir, "requirements.txt") - file, err := os.Open(path) - if errors.Is(err, os.ErrNotExist) { - // having a requirements.txt is optional - } else if err != nil { - return nil, err + // Load Lambda configuration from ol.yaml first (needed to check OL_ENTRY_FILE) + lambdaConfig, err := common.LoadLambdaConfig(codeDir) + if err != nil { + return nil, fmt.Errorf("failed to parse lambda configuration file: %v", err) } - defer file.Close() - - scnr := bufio.NewScanner(file) - for scnr.Scan() { - line := strings.ReplaceAll(scnr.Text(), " ", "") - pkg := strings.Split(line, "#")[0] - if pkg != "" { - pkg = packages.NormalizePkg(pkg) - sandboxMeta.Installs = append(sandboxMeta.Installs, pkg) + + // Determine the Python entry file (default to f.py) + pythonEntryFile := "f.py" + if lambdaConfig.Environment != nil { + if entryFile, ok := lambdaConfig.Environment["OL_ENTRY_FILE"]; ok { + pythonEntryFile = entryFile } } - // Load Lambda configuration from ol.yaml - lambdaConfig, err := common.LoadLambdaConfig(codeDir) - if err != nil { - return nil, fmt.Errorf("failed to parse lambda configuration file: %v", err) + // Determine runtime type by checking for entry file or f.bin + // TODO: support OL_ENTRY_FILE for native runtime + if _, err := os.Stat(filepath.Join(codeDir, pythonEntryFile)); err == nil { + sandboxMeta.Runtime = common.RT_PYTHON + } else if _, err := os.Stat(filepath.Join(codeDir, "f.bin")); err == nil { + sandboxMeta.Runtime = common.RT_NATIVE + } else { + return nil, fmt.Errorf("cannot determine runtime: no %s or f.bin found in %s", pythonEntryFile, codeDir) + } + + // Parse requirements.txt for Python functions (optional) + if sandboxMeta.Runtime == common.RT_PYTHON { + path := filepath.Join(codeDir, "requirements.txt") + file, err := os.Open(path) + if err != nil && !errors.Is(err, os.ErrNotExist) { + return nil, err + } + if err == nil { + defer file.Close() + scnr := bufio.NewScanner(file) + for scnr.Scan() { + line := strings.ReplaceAll(scnr.Text(), " ", "") + pkg := strings.Split(line, "#")[0] + if pkg != "" { + pkg = packages.NormalizePkg(pkg) + sandboxMeta.Installs = append(sandboxMeta.Installs, pkg) + } + } + } } // Return combined FunctionMeta @@ -125,7 +148,7 @@ func (f *LambdaFunc) pullHandlerIfStale() (err error) { } // is there new code? - rtType, codeDir, err := f.lmgr.HandlerPuller.Pull(f.name) + codeDir, err := f.lmgr.HandlerPuller.Pull(f.name) if err != nil { return err } @@ -134,7 +157,11 @@ func (f *LambdaFunc) pullHandlerIfStale() (err error) { return nil } - f.rtType = rtType + // Parse meta to get runtime type and config + meta, err := parseMeta(codeDir) + if err != nil { + return err + } defer func() { if err != nil { @@ -142,7 +169,7 @@ func (f *LambdaFunc) pullHandlerIfStale() (err error) { slog.Error(fmt.Sprintf("could not cleanup %s after failed pull", codeDir)) } - if rtType == common.RT_PYTHON { + if meta.Sandbox.Runtime == common.RT_PYTHON { // we dirty this dir (e.g., by setting up // symlinks to packages, so we want the // HandlerPuller to give us a new one next @@ -152,14 +179,7 @@ func (f *LambdaFunc) pullHandlerIfStale() (err error) { } }() - if rtType == common.RT_PYTHON { - // inspect new code for dependencies; if we can install - // everything necessary, start using new code - meta, err := parseMeta(codeDir) - if err != nil { - return err - } - + if meta.Sandbox.Runtime == common.RT_PYTHON { // make sure all specified dependencies are installed // (but don't recursively find others) for _, pkg := range meta.Sandbox.Installs { @@ -169,17 +189,32 @@ func (f *LambdaFunc) pullHandlerIfStale() (err error) { } f.lmgr.DepTracer.TraceFunction(codeDir, meta.Sandbox.Installs) - f.Meta = meta - } else if rtType == common.RT_NATIVE { + } else if meta.Sandbox.Runtime == common.RT_NATIVE { slog.Info("Got native function") + } - // Initialize f.Meta for native functions for consistensy. - f.Meta = &FunctionMeta{ - Sandbox: nil, // Sandbox is nil for native functions - Config: common.LoadDefaultLambdaConfig(), // Load default configuration + // Write environment variables to .env file if any are specified + if meta.Config.Environment != nil && len(meta.Config.Environment) > 0 { + slog.Info("creating .env for lambda", "entries", len(meta.Config.Environment)) + envPath := filepath.Join(codeDir, ".env") + envFile, err := os.Create(envPath) + if err != nil { + return fmt.Errorf("failed to create .env file: %w", err) + } + defer envFile.Close() + + for key, value := range meta.Config.Environment { + // Quote the value if it contains spaces or special characters + if strings.ContainsAny(value, " \t\n#=") { + escapedValue := strings.ReplaceAll(value, `"`, `\"`) + fmt.Fprintf(envFile, "%s=\"%s\"\n", key, escapedValue) + } else { + fmt.Fprintf(envFile, "%s=%s\n", key, value) + } } } + f.Meta = meta f.codeDir = codeDir f.lastPull = &now return nil diff --git a/go/worker/lambda/lambdaInstance.go b/go/worker/lambda/lambdaInstance.go index af4300ea9..b1e4c3783 100644 --- a/go/worker/lambda/lambdaInstance.go +++ b/go/worker/lambda/lambdaInstance.go @@ -78,6 +78,8 @@ func (linst *LambdaInstance) Task() { return } + reuse := linst.meta.Config.ReuseSandbox + t := common.T0("LambdaInstance-WaitSandbox") // if we have a sandbox, try unpausing it to see if it is still alive if sb != nil { @@ -99,11 +101,11 @@ func (linst *LambdaInstance) Task() { if sb == nil { sb = nil - if f.lmgr.ZygoteProvider != nil && f.rtType == common.RT_PYTHON { + if f.lmgr.ZygoteProvider != nil && linst.meta.Sandbox.Runtime == common.RT_PYTHON { scratchDir := f.lmgr.scratchDirs.Make(f.name) // we don't specify parent SB, because ImportCache.Create chooses it for us - sb, err = f.lmgr.ZygoteProvider.Create(f.lmgr.sbPool, true, linst.codeDir, scratchDir, linst.meta.Sandbox, f.rtType) + sb, err = f.lmgr.ZygoteProvider.Create(f.lmgr.sbPool, true, linst.codeDir, scratchDir, linst.meta.Sandbox) if err != nil { f.printf("failed to get Sandbox from import cache") sb = nil @@ -116,7 +118,7 @@ func (linst *LambdaInstance) Task() { if sb == nil { t2 := common.T0("LambdaInstance-WaitSandbox-NoImportCache") scratchDir := f.lmgr.scratchDirs.Make(f.name) - sb, err = f.lmgr.sbPool.Create(nil, true, linst.codeDir, scratchDir, linst.meta.Sandbox, f.rtType) + sb, err = f.lmgr.sbPool.Create(nil, true, linst.codeDir, scratchDir, linst.meta.Sandbox) t2.T1() } @@ -143,11 +145,22 @@ func (linst *LambdaInstance) Task() { if err != nil { linst.TrySendError(req, http.StatusInternalServerError, "Could not create NewRequest: "+err.Error(), sb) } else { + // Copy headers from original request + for k, vv := range req.r.Header { + for _, v := range vv { + httpReq.Header.Add(k, v) + } + } + // Preserve ContentLength (parsed from Content-Length header) + httpReq.ContentLength = req.r.ContentLength + resp, err := sb.Client().Do(httpReq) // copy response out if err != nil { linst.TrySendError(req, http.StatusBadGateway, "RoundTrip failed: "+err.Error()+"\n", sb) + sb.Destroy("Sandbox's HTTP client returned an error") + sb = nil } else { // copy headers // (adapted from copyHeaders: https://go.dev/src/net/http/httputil/reverseproxy.go) @@ -181,20 +194,27 @@ func (linst *LambdaInstance) Task() { } f.doneChan <- req + // If reuse is disabled, destroy the sandbox after invocation. + if !reuse && sb != nil { + sb.Destroy("reuse-sandbox disabled: destroying sandbox after invocation") + sb = nil + } + // check whether we should shutdown (non-blocking) select { case killed := <-linst.killChan: - rtLog := sb.GetRuntimeLog() - sb.Destroy("Lambda instance kill signal received") - - slog.Info("Stopped sandbox") - - if common.Conf.Log_output { - if rtLog != "" { - slog.Info("Runtime output is:") - - for _, line := range strings.Split(rtLog, "\n") { - slog.Info(fmt.Sprintf(" %s", line)) + if sb != nil { + rtLog := sb.GetRuntimeLog() + sb.Destroy("Lambda instance kill signal received") + slog.Info("Stopped sandbox") + + if common.Conf.Log_output { + if rtLog != "" { + slog.Info("Runtime output is:") + + for _, line := range strings.Split(rtLog, "\n") { + slog.Info(fmt.Sprintf(" %s", line)) + } } } } @@ -210,11 +230,19 @@ func (linst *LambdaInstance) Task() { default: req = nil } + + // if sandbox was destroyed, break out so outer loop can create a new one + if sb == nil { + break + } } if sb != nil { + if !reuse { + panic("sb should be nil when reuse is disabled") + } if err := sb.Pause(); err != nil { - f.printf("discard sandbox %s due to Pause error: %v", sb.ID(), err) + f.printf("Discard sandbox %s due to Pause error: %v", sb.ID(), err) sb = nil } } diff --git a/go/worker/lambda/packages/packagePuller.go b/go/worker/lambda/packages/packagePuller.go index 1546e4cad..d971ee77b 100644 --- a/go/worker/lambda/packages/packagePuller.go +++ b/go/worker/lambda/packages/packagePuller.go @@ -180,9 +180,10 @@ func (pp *PackagePuller) sandboxInstall(p *Package) (err error) { inst := common.Conf.InstallerLimits.WithDefaults(&common.Conf.Limits) meta := &sandbox.SandboxMeta{ + Runtime: common.RT_PYTHON, MemLimitMB: inst.Mem_mb, } - sb, err := pp.sbPool.Create(nil, true, pp.pipLambda, scratchDir, meta, common.RT_PYTHON) + sb, err := pp.sbPool.Create(nil, true, pp.pipLambda, scratchDir, meta) if err != nil { return err } diff --git a/go/worker/lambda/zygote/api.go b/go/worker/lambda/zygote/api.go index 36b721878..52d666d4f 100644 --- a/go/worker/lambda/zygote/api.go +++ b/go/worker/lambda/zygote/api.go @@ -1,13 +1,11 @@ package zygote import ( - "github.com/open-lambda/open-lambda/go/common" "github.com/open-lambda/open-lambda/go/worker/sandbox" ) type ZygoteProvider interface { Create(childSandboxPool sandbox.SandboxPool, isLeaf bool, - codeDir, scratchDir string, meta *sandbox.SandboxMeta, - rt_type common.RuntimeType) (sandbox.Sandbox, error) + codeDir, scratchDir string, meta *sandbox.SandboxMeta) (sandbox.Sandbox, error) Cleanup() } diff --git a/go/worker/lambda/zygote/importCache.go b/go/worker/lambda/zygote/importCache.go index a0145be98..743daf981 100644 --- a/go/worker/lambda/zygote/importCache.go +++ b/go/worker/lambda/zygote/importCache.go @@ -146,7 +146,7 @@ func (cache *ImportCache) recursiveKill(node *ImportCacheNode) { } // Create creates a new sandbox using the import cache. -func (cache *ImportCache) Create(childSandboxPool sandbox.SandboxPool, isLeaf bool, codeDir, scratchDir string, meta *sandbox.SandboxMeta, rt_type common.RuntimeType) (sandbox.Sandbox, error) { +func (cache *ImportCache) Create(childSandboxPool sandbox.SandboxPool, isLeaf bool, codeDir, scratchDir string, meta *sandbox.SandboxMeta) (sandbox.Sandbox, error) { t := common.T0("ImportCache.Create") defer t.T1() @@ -158,7 +158,7 @@ func (cache *ImportCache) Create(childSandboxPool sandbox.SandboxPool, isLeaf bo panic(fmt.Errorf("did not find Zygote; at least expected to find the root")) } slog.Info(fmt.Sprintf("Try using Zygote from <%v>", node)) - return cache.createChildSandboxFromNode(childSandboxPool, node, isLeaf, codeDir, scratchDir, meta, rt_type) + return cache.createChildSandboxFromNode(childSandboxPool, node, isLeaf, codeDir, scratchDir, meta) } // use getSandboxInNode to get a Zygote Sandbox for the node (creating one @@ -167,20 +167,20 @@ func (cache *ImportCache) Create(childSandboxPool sandbox.SandboxPool, isLeaf bo // the new Sandbox may either be for a Zygote, or a leaf Sandbox func (cache *ImportCache) createChildSandboxFromNode( childSandboxPool sandbox.SandboxPool, node *ImportCacheNode, isLeaf bool, - codeDir, scratchDir string, meta *sandbox.SandboxMeta, rt_type common.RuntimeType) (sandbox.Sandbox, error) { + codeDir, scratchDir string, meta *sandbox.SandboxMeta) (sandbox.Sandbox, error) { t := common.T0("ImportCache.createChildSandboxFromNode") defer t.T1() // try twice, restarting parent Sandbox if it fails the first time forceNew := false for i := 0; i < 2; i++ { - zygoteSB, isNew, err := cache.getSandboxInNode(node, forceNew, rt_type) + zygoteSB, isNew, err := cache.getSandboxInNode(node, forceNew) if err != nil { return nil, err } t2 := common.T0("ImportCache.createChildSandboxFromNode:childSandboxPool.Create") - sb, err := childSandboxPool.Create(zygoteSB, isLeaf, codeDir, scratchDir, meta, rt_type) + sb, err := childSandboxPool.Create(zygoteSB, isLeaf, codeDir, scratchDir, meta) if err == nil { if isLeaf { @@ -212,7 +212,7 @@ func (cache *ImportCache) createChildSandboxFromNode( // // the Sandbox returned is guaranteed to be in Unpaused state. After // use, caller must also call putSandboxInNode to release ref count -func (cache *ImportCache) getSandboxInNode(node *ImportCacheNode, forceNew bool, rt_type common.RuntimeType) (sb sandbox.Sandbox, isNew bool, err error) { +func (cache *ImportCache) getSandboxInNode(node *ImportCacheNode, forceNew bool) (sb sandbox.Sandbox, isNew bool, err error) { t := common.T0("ImportCache.getSandboxInNode") defer t.T1() @@ -239,7 +239,7 @@ func (cache *ImportCache) getSandboxInNode(node *ImportCacheNode, forceNew bool, } // SLOW PATH - if err := cache.createSandboxInNode(node, rt_type); err != nil { + if err := cache.createSandboxInNode(node); err != nil { return nil, false, err } node.sbRefCount = 1 @@ -281,7 +281,7 @@ func (*ImportCache) putSandboxInNode(node *ImportCacheNode, sb sandbox.Sandbox) } } -func (cache *ImportCache) createSandboxInNode(node *ImportCacheNode, rt_type common.RuntimeType) (err error) { +func (cache *ImportCache) createSandboxInNode(node *ImportCacheNode) (err error) { // populate codeDir/packages with deps, and record top-level mods) if node.codeDir == "" { codeDir := cache.codeDirs.Make("import-cache") @@ -305,7 +305,9 @@ func (cache *ImportCache) createSandboxInNode(node *ImportCacheNode, rt_type com // policy: what modules should we pre-import? Top-level of // pre-initialized packages is just one possibility... + // Import cache is Python-specific, so always use RT_PYTHON node.meta = &sandbox.SandboxMeta{ + Runtime: common.RT_PYTHON, Installs: installs, Imports: topLevelMods, } @@ -314,9 +316,9 @@ func (cache *ImportCache) createSandboxInNode(node *ImportCacheNode, rt_type com scratchDir := cache.scratchDirs.Make("import-cache") var sb sandbox.Sandbox if node.parent != nil { - sb, err = cache.createChildSandboxFromNode(cache.sbPool, node.parent, false, node.codeDir, scratchDir, node.meta, rt_type) + sb, err = cache.createChildSandboxFromNode(cache.sbPool, node.parent, false, node.codeDir, scratchDir, node.meta) } else { - sb, err = cache.sbPool.Create(nil, false, node.codeDir, scratchDir, node.meta, common.RT_PYTHON) + sb, err = cache.sbPool.Create(nil, false, node.codeDir, scratchDir, node.meta) } if err != nil { diff --git a/go/worker/lambda/zygote/multiTree.go b/go/worker/lambda/zygote/multiTree.go index e0a29637a..aa607bef5 100644 --- a/go/worker/lambda/zygote/multiTree.go +++ b/go/worker/lambda/zygote/multiTree.go @@ -44,9 +44,9 @@ func NewMultiTree(codeDirs *common.DirMaker, scratchDirs *common.DirMaker, sbPoo } // Create creates a new sandbox using a randomly selected ImportCache tree. -func (mt *MultiTree) Create(childSandboxPool sandbox.SandboxPool, isLeaf bool, codeDir, scratchDir string, meta *sandbox.SandboxMeta, rt_type common.RuntimeType) (sandbox.Sandbox, error) { +func (mt *MultiTree) Create(childSandboxPool sandbox.SandboxPool, isLeaf bool, codeDir, scratchDir string, meta *sandbox.SandboxMeta) (sandbox.Sandbox, error) { idx := rand.Intn(len(mt.trees)) - return mt.trees[idx].Create(childSandboxPool, isLeaf, codeDir, scratchDir, meta, rt_type) + return mt.trees[idx].Create(childSandboxPool, isLeaf, codeDir, scratchDir, meta) } // Cleanup performs cleanup operations for all ImportCache trees in the MultiTree. diff --git a/go/worker/sandbox/api.go b/go/worker/sandbox/api.go index bbe215cf3..b109e9496 100644 --- a/go/worker/sandbox/api.go +++ b/go/worker/sandbox/api.go @@ -12,8 +12,8 @@ type SandboxPool interface { // isLeaf: true iff this is not being created as a sandbox we can fork later // codeDir: directory where lambda code exists // scratchDir: directory where handler code can write (caller is responsible for creating and deleting) - // meta: details about installs, imports, etc. Will be populated with defaults if not specified - Create(parent Sandbox, isLeaf bool, codeDir, scratchDir string, meta *SandboxMeta, rtType common.RuntimeType) (sb Sandbox, err error) + // meta: details about runtime, installs, imports, etc. Will be populated with defaults if not specified + Create(parent Sandbox, isLeaf bool, codeDir, scratchDir string, meta *SandboxMeta) (sb Sandbox, err error) // blocks until all Sandboxes are deleted, so caller must // either delete them before this call, or from another asyncronously @@ -80,15 +80,16 @@ type Sandbox interface { // Child calls this on parent to notify of child Destroy childExit(child Sandbox) - - GetRuntimeType() common.RuntimeType // TODO: make it part of SandboxMeta? } type SandboxMeta struct { - Installs []string - Imports []string + Runtime common.RuntimeType MemLimitMB int CPUPercent int + + // Python specific fields: + Installs []string + Imports []string } type SandboxError string diff --git a/go/worker/sandbox/cgroups/api.go b/go/worker/sandbox/cgroups/api.go index ae41c7ba8..69fe3c6e7 100644 --- a/go/worker/sandbox/cgroups/api.go +++ b/go/worker/sandbox/cgroups/api.go @@ -7,10 +7,9 @@ type Cgroup interface { SetMemLimitMB(mb int) Pause() error Unpause() error - Release() AddPid(pid string) error GetPIDs() ([]string, error) - KillAllProcs() + KillAndRelease() DebugString() string // TODO: find a way to rip this out. Higher layers should not diff --git a/go/worker/sandbox/cgroups/cgroup.go b/go/worker/sandbox/cgroups/cgroup.go index f5d2fddf6..35d18bf3a 100644 --- a/go/worker/sandbox/cgroups/cgroup.go +++ b/go/worker/sandbox/cgroups/cgroup.go @@ -2,7 +2,10 @@ package cgroups import ( "bufio" + "errors" "fmt" + "golang.org/x/sys/unix" + "io" "io/ioutil" "log/slog" "os" @@ -32,27 +35,18 @@ func (cg *CgroupImpl) Name() string { return cg.name } -// Release releases the cgroup back to the pool or destroys it if the pool is full. -func (cg *CgroupImpl) Release() { +// KillAndRelease stops all processes inside the cgroup. +// After releasing, the cgroup can be recycled or destroyed depending on configuration. +// Note, the CG most be paused beforehand. +func (cg *CgroupImpl) KillAndRelease() { + err := cg.WriteEventAndWait("cgroup.kill", 1, "populated", 0, 20*time.Second) + if err != nil { + panic(fmt.Errorf("can't write \"1\" to cgroup.kill: %w", err)) + } + // if there's room in the recycled channel, add it there. // Otherwise, just delete it. if common.Conf.Features.Reuse_cgroups { - for i := 100; i >= 0; i-- { - pids, err := cg.GetPIDs() - if err != nil { - panic(err) - } else if len(pids) > 0 { - if i == 0 { - panic(fmt.Errorf("Cannot release cgroup that contains processes: %v", pids)) - } - - cg.printf("cgroup Rmdir failed, trying again in 5ms") - time.Sleep(5 * time.Millisecond) - } else { - break - } - } - select { case cg.pool.recycled <- cg: cg.printf("release and recycle") @@ -84,9 +78,9 @@ func (cg *CgroupImpl) Destroy() { } } -// GroupPath returns the path to the Cgroup pool for OpenLambda +// GroupPath returns the path to this cgroup directory. func (cg *CgroupImpl) GroupPath() string { - return fmt.Sprintf("%s/%s", cg.pool.GroupPath(), cg.name) + return fmt.Sprintf("%s/%s", cg.pool.poolPath, cg.name) } func (cg *CgroupImpl) MemoryEvents() map[string]int64 { @@ -113,7 +107,7 @@ func (cg *CgroupImpl) MemoryEvents() map[string]int64 { // ResourcePath returns the path to a specific resource in this cgroup func (cg *CgroupImpl) ResourcePath(resource string) string { - return fmt.Sprintf("%s/%s/%s", cg.pool.GroupPath(), cg.name, resource) + return fmt.Sprintf("%s/%s/%s", cg.pool.poolPath, cg.name, resource) } func (cg *CgroupImpl) TryWriteInt(resource string, val int64) error { @@ -130,18 +124,81 @@ func (cg *CgroupImpl) WriteInt(resource string, val int64) { } } +// WriteEventAndWait() writes to cgroup controller file and waits for the corresponding event in cgroup.events to be updated +func (cg *CgroupImpl) WriteEventAndWait(controller string, controllerState int64, event string, eventState int64, timeout time.Duration) error { + resourcePath := cg.ResourcePath("cgroup.events") + eventFile, err := os.Open(resourcePath) + if err != nil { + return fmt.Errorf("failed to open %s: %w", resourcePath, err) + } + + // cgroups(7): POLLPRI indicates "cgroup.events file modified" + // for poll to decide a POLLPRI event occurs it maintains 2 event counters: + // 1. the event counter when you last read the file + // 2. the file's current event counter + // if the last read's counter is different from the current event counter poll returns POLLPRI + pollFDs := []unix.PollFd{ + { + Fd: int32(eventFile.Fd()), + Events: unix.POLLPRI, + }, + } + pollCalls := 0 + + start := time.Now() + + defer func() { + elapsed := time.Since(start) + if elapsed >= 250*time.Millisecond { + cg.printf("WARNING! WriteEventAndWait to state %v took %v to complete", controllerState, elapsed) + } + if pollCalls > 5 { + cg.printf("WARNING! WriteEventAndWait called poll %v times, could be busy waiting", pollCalls) + } + }() + + cg.WriteInt(controller, controllerState) + for { + elapsed := time.Since(start) + + remaining := timeout - elapsed + if remaining < 0 { + return fmt.Errorf("%s timeout after %v (expected state %v)", controller, timeout, eventState) + } + + pollCalls++ + _, err := unix.Poll(pollFDs, int(remaining.Milliseconds())) + if err != nil && !errors.Is(err, unix.EINTR) { + return fmt.Errorf("poll syscall failed on %s: %w", resourcePath, err) + } + + // read from the same file to update event counter, prevents busy wait + currEventState, err := cg.TryReadIntKVFromFile(eventFile, event) + if err != nil { + return fmt.Errorf("failed to check %s in %s :: %w", event, resourcePath, err) + } + if currEventState == eventState { + return nil + } + } +} + func (cg *CgroupImpl) WriteString(resource string, val string) { if err := cg.TryWriteString(resource, val); err != nil { panic(fmt.Sprintf("Error writing %v to %s: %v", val, resource, err)) } } -func (cg *CgroupImpl) TryReadIntKV(resource string, key string) (int64, error) { - raw, err := ioutil.ReadFile(cg.ResourcePath(resource)) +func (_ *CgroupImpl) TryReadIntKVFromFile(file *os.File, key string) (int64, error) { + _, err := file.Seek(0, io.SeekStart) if err != nil { - return 0, err + return 0, fmt.Errorf("failed to seek to start of file: %w", err) + } + data, err := io.ReadAll(file) + if err != nil { + return 0, fmt.Errorf("failed to read key %s from file: %w", key, err) } - body := string(raw) + body := string(data) lines := strings.Split(body, "\n") for i := 0; i <= len(lines); i++ { parts := strings.Split(lines[i], " ") @@ -156,6 +213,16 @@ func (cg *CgroupImpl) TryReadIntKV(resource string, key string) (int64, error) { return 0, fmt.Errorf("could not find key '%s' in file: %s", key, body) } +func (cg *CgroupImpl) TryReadIntKV(resource string, key string) (int64, error) { + resourcePath := cg.ResourcePath(resource) + file, err := os.Open(resourcePath) + if err != nil { + return 0, fmt.Errorf("failed to open file %s: %w", resourcePath, err) + } + defer file.Close() + return cg.TryReadIntKVFromFile(file, key) +} + func (cg *CgroupImpl) TryReadInt(resource string) (int64, error) { raw, err := ioutil.ReadFile(cg.ResourcePath(resource)) if err != nil { @@ -189,27 +256,8 @@ func (cg *CgroupImpl) AddPid(pid string) error { } func (cg *CgroupImpl) setFreezeState(state int64) error { - cg.WriteInt("cgroup.freeze", state) - - timeout := 5 * time.Second - - start := time.Now() - for { - freezerState, err := cg.TryReadInt("cgroup.freeze") - if err != nil { - return fmt.Errorf("failed to check self_freezing state :: %v", err) - } - - if freezerState == state { - return nil - } - - if time.Since(start) > timeout { - return fmt.Errorf("cgroup stuck on %v after %v (should be %v)", freezerState, timeout, state) - } - - time.Sleep(1 * time.Millisecond) - } + timeout := 20 * time.Second + return cg.WriteEventAndWait("cgroup.freeze", state, "frozen", state, timeout) } // get mem usage in MB @@ -298,12 +346,6 @@ func (cg *CgroupImpl) CgroupProcsPath() string { return cg.ResourcePath("cgroup.procs") } -// KillAllProcs stops all processes inside the cgroup. -// Note, the CG most be paused beforehand -func (cg *CgroupImpl) KillAllProcs() { - cg.WriteInt("cgroup.kill", 1) -} - // DebugString returns a string representation of the cgroup's state. func (cg *CgroupImpl) DebugString() string { s := "" diff --git a/go/worker/sandbox/cgroups/pool.go b/go/worker/sandbox/cgroups/pool.go index 13383fff5..7d83a3b43 100644 --- a/go/worker/sandbox/cgroups/pool.go +++ b/go/worker/sandbox/cgroups/pool.go @@ -2,13 +2,12 @@ package cgroups import ( "fmt" - "io/ioutil" "log/slog" "os" - "path" + "path/filepath" + "strconv" "strings" "syscall" - "time" "github.com/open-lambda/open-lambda/go/common" ) @@ -19,35 +18,56 @@ const CGROUP_RESERVE = 16 type CgroupPool struct { Name string + poolPath string ready chan *CgroupImpl recycled chan *CgroupImpl quit chan chan bool nextID int } -// NewCgroupPool creates a new CgroupPool with the specified name. -func NewCgroupPool(name string) (*CgroupPool, error) { +// InitPoolRoot creates the cgroup pool root directory and enables controllers. +func InitPoolRoot(poolPath string) error { + + if err := os.MkdirAll(poolPath, 0700); err != nil { + return fmt.Errorf("failed to create cgroup pool root %s: %w", poolPath, err) + } + + ctrlPath := filepath.Join(poolPath, "cgroup.subtree_control") + if err := os.WriteFile(ctrlPath, []byte("+pids +io +memory +cpu"), os.ModeAppend); err != nil { + return fmt.Errorf("failed to enable controllers at %s: %w", ctrlPath, err) + } + + uidStr := os.Getenv("SUDO_UID") + if uidStr == "" { + return fmt.Errorf("SUDO_UID not set; worker must be run with sudo") + } + uid, err := strconv.Atoi(uidStr) + if err != nil { + return fmt.Errorf("invalid SUDO_UID value %q: %w", uidStr, err) + } + if err := os.Chown(poolPath, uid, uid); err != nil { + return fmt.Errorf("failed to chown cgroup pool root: %w", err) + } + + fmt.Printf("\tCreated cgroup pool root at %s\n", poolPath) + return nil +} + +func NewCgroupPool(name string, poolPath string) (*CgroupPool, error) { pool := &CgroupPool{ - Name: path.Base(path.Dir(common.Conf.Worker_dir)) + "-" + name, + Name: name, + poolPath: poolPath, ready: make(chan *CgroupImpl, CGROUP_RESERVE), recycled: make(chan *CgroupImpl, CGROUP_RESERVE), quit: make(chan chan bool), nextID: 0, } - // create cgroup - groupPath := pool.GroupPath() - pool.printf("create %s", groupPath) - if err := syscall.Mkdir(groupPath, 0700); err != nil { - return nil, fmt.Errorf("Mkdir %s: %s", groupPath, err) - } - - // Make controllers available to child groups - rpath := fmt.Sprintf("%s/cgroup.subtree_control", groupPath) - if err := ioutil.WriteFile(rpath, []byte("+pids +io +memory +cpu"), os.ModeAppend); err != nil { - panic(fmt.Sprintf("Error writing to %s: %v", rpath, err)) + if st, err := os.Stat(poolPath); err != nil || !st.IsDir() { + return nil, fmt.Errorf("cgroup pool root %s does not exist.", poolPath) } + pool.printf("reusing pool root %s", poolPath) go pool.cgTask() return pool, nil } @@ -136,28 +156,14 @@ Empty: done <- true } -// Destroy this entire cgroup pool +// Destroy drains all child cgroups but preserves the pool root. func (pool *CgroupPool) Destroy() { // signal cgTask, then wait for it to finish ch := make(chan bool) pool.quit <- ch <-ch - // Destroy cgroup for this entire pool - gpath := pool.GroupPath() - pool.printf("Destroying cgroup pool with path \"%s\"", gpath) - for i := 100; i >= 0; i-- { - if err := syscall.Rmdir(gpath); err != nil { - if i == 0 { - panic(fmt.Errorf("Rmdir %s: %s", gpath, err)) - } - - pool.printf("cgroup pool Rmdir failed, trying again in 5ms") - time.Sleep(5 * time.Millisecond) - } else { - break - } - } + pool.printf("destroyed all child cgroups, pool root preserved") } // GetCg retrieves a cgroup from the pool, setting its memory limit and CPU percentage. @@ -178,8 +184,3 @@ func (pool *CgroupPool) GetCg(memLimitMB int, moveMemCharge bool, cpuPercent int return cg } - -// GroupPath returns the path to the Cgroup pool for OpenLambda -func (pool *CgroupPool) GroupPath() string { - return fmt.Sprintf("/sys/fs/cgroup/%s", pool.Name) -} diff --git a/go/worker/sandbox/docker.go b/go/worker/sandbox/docker.go index 5087f5d81..80ea35ef6 100644 --- a/go/worker/sandbox/docker.go +++ b/go/worker/sandbox/docker.go @@ -33,7 +33,6 @@ type DockerContainer struct { client *docker.Client installed map[string]bool meta *SandboxMeta - rtType common.RuntimeType httpClient *http.Client } @@ -253,11 +252,6 @@ func (container *DockerContainer) ID() string { return container.hostID } -// GetRuntimeType returns what runtime is being used by this container? -func (container *DockerContainer) GetRuntimeType() common.RuntimeType { - return container.rtType -} - // DockerID returns the id assigned by docker itself, not by open lambda func (container *DockerContainer) DockerID() string { return container.container.ID @@ -269,7 +263,7 @@ func (container *DockerContainer) HostDir() string { } func (container *DockerContainer) runServer() error { - if container.rtType != common.RT_PYTHON { + if container.meta.Runtime != common.RT_PYTHON { return fmt.Errorf("Unsupported runtime") } diff --git a/go/worker/sandbox/dockerPool.go b/go/worker/sandbox/dockerPool.go index f00bd87a2..28dd30d3f 100644 --- a/go/worker/sandbox/dockerPool.go +++ b/go/worker/sandbox/dockerPool.go @@ -60,7 +60,7 @@ func NewDockerPool(pidMode string, caps []string) (*DockerPool, error) { } // Create creates a docker sandbox from the handler and sandbox directory. -func (pool *DockerPool) Create(parent Sandbox, isLeaf bool, codeDir, scratchDir string, meta *SandboxMeta, _ common.RuntimeType) (sb Sandbox, err error) { +func (pool *DockerPool) Create(parent Sandbox, isLeaf bool, codeDir, scratchDir string, meta *SandboxMeta) (sb Sandbox, err error) { meta = fillMetaDefaults(meta) t := common.T0("Create()") defer t.T1() diff --git a/go/worker/sandbox/mock.go b/go/worker/sandbox/mock.go new file mode 100644 index 000000000..b4cbf4d66 --- /dev/null +++ b/go/worker/sandbox/mock.go @@ -0,0 +1,122 @@ +package sandbox + +import ( + "fmt" + "net/http" + "sync" + "sync/atomic" +) + +// MockSandbox is a test double for Sandbox. +// Exported fields control error injection; state fields track lifecycle. +type MockSandbox struct { + mu sync.Mutex + id string + paused bool + destroyed bool + + // Set these before calling Get/Put to inject errors. + PauseErr error + UnpauseErr error +} + +var mockIDCounter int64 + +// NewMockSandbox creates a MockSandbox with the given ID. +func NewMockSandbox(id string) *MockSandbox { + return &MockSandbox{id: id, paused: true} +} + +func (m *MockSandbox) ID() string { return m.id } + +func (m *MockSandbox) Destroy(reason string) { + m.mu.Lock() + defer m.mu.Unlock() + m.destroyed = true +} + +func (m *MockSandbox) DestroyIfPaused(reason string) { + m.mu.Lock() + defer m.mu.Unlock() + if m.paused { + m.destroyed = true + } +} + +func (m *MockSandbox) Pause() error { + m.mu.Lock() + defer m.mu.Unlock() + if m.PauseErr != nil { + return m.PauseErr + } + m.paused = true + return nil +} + +func (m *MockSandbox) Unpause() error { + m.mu.Lock() + defer m.mu.Unlock() + if m.UnpauseErr != nil { + return m.UnpauseErr + } + m.paused = false + return nil +} + +func (m *MockSandbox) Client() *http.Client { return nil } +func (m *MockSandbox) Meta() *SandboxMeta { return nil } +func (m *MockSandbox) GetRuntimeLog() string { return "" } +func (m *MockSandbox) GetProxyLog() string { return "" } +func (m *MockSandbox) DebugString() string { return fmt.Sprintf("mock:%s", m.id) } +func (m *MockSandbox) fork(dst Sandbox) error { return nil } +func (m *MockSandbox) childExit(child Sandbox) {} + +// IsDestroyed returns whether Destroy has been called. +func (m *MockSandbox) IsDestroyed() bool { + m.mu.Lock() + defer m.mu.Unlock() + return m.destroyed +} + +// IsPaused returns the current pause state. +func (m *MockSandbox) IsPaused() bool { + m.mu.Lock() + defer m.mu.Unlock() + return m.paused +} + +// MockSandboxPool is a test double for SandboxPool. +// It creates MockSandbox instances with auto-incremented IDs. +type MockSandboxPool struct { + mu sync.Mutex + Created []*MockSandbox + + // Set before calling Get to make pool.Create fail. + CreateErr error +} + +func (p *MockSandboxPool) Create(parent Sandbox, isLeaf bool, codeDir, scratchDir string, meta *SandboxMeta) (Sandbox, error) { + p.mu.Lock() + defer p.mu.Unlock() + if p.CreateErr != nil { + return nil, p.CreateErr + } + id := fmt.Sprintf("mock-%d", atomic.AddInt64(&mockIDCounter, 1)) + sb := NewMockSandbox(id) + sb.paused = false // Pool.Create returns unpaused sandboxes + p.Created = append(p.Created, sb) + return sb, nil +} + +func (p *MockSandboxPool) Cleanup() {} +func (p *MockSandboxPool) AddListener(handler SandboxEventFunc) {} +func (p *MockSandboxPool) DebugString() string { return "mock-pool" } + +// CreatedSandboxes returns a snapshot of all sandboxes created by this pool. +func (p *MockSandboxPool) CreatedSandboxes() []*MockSandbox { + p.mu.Lock() + defer p.mu.Unlock() + out := make([]*MockSandbox, len(p.Created)) + copy(out, p.Created) + return out +} diff --git a/go/worker/sandbox/sandbox.go b/go/worker/sandbox/sandbox.go index a6a525bf4..2d4ae17ea 100644 --- a/go/worker/sandbox/sandbox.go +++ b/go/worker/sandbox/sandbox.go @@ -18,6 +18,8 @@ func SandboxPoolFromConfig(name string, sizeMb int) (cf SandboxPool, err error) } NewSOCKEvictor(pool) return pool, nil + } else if common.Conf.Sandbox == "mock" { + return &MockSandboxPool{}, nil } return nil, fmt.Errorf("invalid sandbox type: '%s'", common.Conf.Sandbox) diff --git a/go/worker/sandbox/sock.go b/go/worker/sandbox/sock.go index 9ee3c9491..5f813c93b 100644 --- a/go/worker/sandbox/sock.go +++ b/go/worker/sandbox/sock.go @@ -26,7 +26,6 @@ type SOCKContainer struct { codeDir string scratchDir string cg cgroups.Cgroup - rtType common.RuntimeType client *http.Client // 1 for self, plus 1 for each child (we can't release memory @@ -53,10 +52,6 @@ func (container *SOCKContainer) ID() string { return container.id } -func (container *SOCKContainer) GetRuntimeType() common.RuntimeType { - return container.rtType -} - func (container *SOCKContainer) freshProc() (err error) { // get FD to cgroup cgFiles := make([]*os.File, 1) @@ -70,13 +65,13 @@ func (container *SOCKContainer) freshProc() (err error) { var cmd *exec.Cmd - if container.rtType == common.RT_PYTHON { + if container.meta.Runtime == common.RT_PYTHON { cmd = exec.Command( "chroot", container.containerRootDir, "python3", "-u", "/runtimes/python/server.py", "/host/bootstrap.py", strconv.Itoa(1), strconv.FormatBool(common.Conf.Features.Enable_seccomp), ) - } else if container.rtType == common.RT_NATIVE { + } else if container.meta.Runtime == common.RT_NATIVE { if container.containerProxy == nil { err := container.launchContainerProxy() @@ -282,7 +277,7 @@ func (container *SOCKContainer) DestroyIfPaused(reason string) { } // when the count goes to zero, it means (a) this container and (b) -// all it's descendants are destroyed. Thus, it's safe to release it's +// all its descendants are destroyed. Thus, it's safe to release its // cgroups, and return the memory allocation to the memPool func (container *SOCKContainer) decCgRefCount() { newCount := atomic.AddInt32(&container.cgRefCount, -1) @@ -302,9 +297,8 @@ func (container *SOCKContainer) decCgRefCount() { t := common.T0("Destroy()/cleanup-cgroup") if container.cg != nil { - container.cg.KillAllProcs() + container.cg.KillAndRelease() container.printf("killed PIDs in CG\n") - container.cg.Release() container.pool.mem.adjustAvailableMB(container.cg.GetMemLimitMB()) } t.T1() @@ -322,6 +316,12 @@ func (container *SOCKContainer) decCgRefCount() { } t.T1() + // Clean up ol.sock from scratchDir (scratchDir itself may be reused, e.g., for package caching) + sockPath := filepath.Join(container.scratchDir, "ol.sock") + if err := os.Remove(sockPath); err != nil && !os.IsNotExist(err) { + container.printf("remove socket %s failed :: %v\n", sockPath, err) + } + if container.parent != nil { container.parent.childExit(container) } diff --git a/go/worker/sandbox/sockPool.go b/go/worker/sandbox/sockPool.go index d8d793cd9..e52e8c3a3 100644 --- a/go/worker/sandbox/sockPool.go +++ b/go/worker/sandbox/sockPool.go @@ -34,7 +34,8 @@ type SOCKPool struct { // NewSOCKPool creates a SOCKPool. func NewSOCKPool(name string, mem *MemPool) (cf *SOCKPool, err error) { - cgPool, err := cgroups.NewCgroupPool(name) + olPath := filepath.Dir(common.Conf.Worker_dir) + cgPool, err := cgroups.NewCgroupPool(name, common.CgroupPoolPath(olPath)) if err != nil { return nil, err } @@ -64,7 +65,7 @@ func sbStr(sb Sandbox) string { return fmt.Sprintf("", sb.ID()) } -func (pool *SOCKPool) Create(parent Sandbox, isLeaf bool, codeDir, scratchDir string, meta *SandboxMeta, rtType common.RuntimeType) (sb Sandbox, err error) { +func (pool *SOCKPool) Create(parent Sandbox, isLeaf bool, codeDir, scratchDir string, meta *SandboxMeta) (sb Sandbox, err error) { id := fmt.Sprintf("%d", atomic.AddInt64(&nextId, 1)) meta = fillMetaDefaults(meta) pool.printf("<%v>.Create(%v, %v, %v, %v, %v)=%s...", pool.name, sbStr(parent), isLeaf, codeDir, scratchDir, meta, id) @@ -84,7 +85,6 @@ func (pool *SOCKPool) Create(parent Sandbox, isLeaf bool, codeDir, scratchDir st cgRefCount: 1, children: make(map[string]Sandbox), meta: meta, - rtType: rtType, containerProxy: nil, } var c Sandbox = cSock @@ -122,7 +122,7 @@ func (pool *SOCKPool) Create(parent Sandbox, isLeaf bool, codeDir, scratchDir st } t2.T1() - if rtType == common.RT_PYTHON { + if meta.Runtime == common.RT_PYTHON { // add installed packages to the path, and import the modules we'll need var pyCode []string @@ -153,7 +153,7 @@ func (pool *SOCKPool) Create(parent Sandbox, isLeaf bool, codeDir, scratchDir st if err := ioutil.WriteFile(path, code, 0600); err != nil { return nil, err } - } else if rtType == common.RT_NATIVE { + } else if meta.Runtime == common.RT_NATIVE { // nothing to do? } else { return nil, fmt.Errorf("Unsupported runtime") diff --git a/go/worker/sandboxset/api.go b/go/worker/sandboxset/api.go new file mode 100644 index 000000000..f78d22791 --- /dev/null +++ b/go/worker/sandboxset/api.go @@ -0,0 +1,67 @@ +// Package sandboxset provides a thread-safe pool of sandboxes for a single Lambda function. Callers ask for a sandbox and don't worry about whether it is freshly created or recycled from a previous request. +// +// Usage: +// +// set := sandboxset.New(&sandboxset.Config{ +// Pool: myPool, +// CodeDir: "/path/to/lambda", +// ScratchDirs: myScratchDirs, +// }) +// +// ref, err := set.GetOrCreateUnpaused() +// // ... use ref.Sandbox() to handle request ... +// if broken { +// ref.MarkDead() +// } +// ref.Put() +package sandboxset + +import ( + "github.com/open-lambda/open-lambda/go/common" + "github.com/open-lambda/open-lambda/go/worker/sandbox" +) + +// SandboxSet manages a pool of sandboxes for one Lambda function. +// All methods are safe to call from multiple goroutines. +type SandboxSet interface { + // GetOrCreateUnpaused returns an unpaused sandbox ready to handle a + // request, wrapped in a SandboxRef. + GetOrCreateUnpaused() (*SandboxRef, error) + + // Close marks the set closed and destroys idle sandboxes. In-use refs + // are destroyed when their holder returns them via Put. + Close() error +} + +// Config holds the parameters needed to create a SandboxSet. +type Config struct { + // Pool creates and destroys the underlying sandboxes. + Pool sandbox.SandboxPool + + // Parent is an optional SandboxSet to fork from. When nil, new + // sandboxes are created from scratch. Not all SandboxPool + // implementations support forking. The parent must outlive this child. + Parent SandboxSet + + // IsLeaf marks sandboxes as non-forkable, meaning they will + // not be used as parents for future forks. + IsLeaf bool + + // CodeDir is the directory containing the Lambda handler code. + CodeDir string + + // Meta holds runtime configuration (memory limits, packages, + // imports, etc.). Nil means the pool fills in defaults. + Meta *sandbox.SandboxMeta + + // ScratchDirs creates a unique writable directory for each + // new sandbox. The set calls ScratchDirs.Make internally + // so that GetOrCreateUnpaused can remain argument-free. + ScratchDirs *common.DirMaker +} + +// New creates a SandboxSet from cfg. +// Panics if Pool, CodeDir, or ScratchDirs are missing. +func New(cfg *Config) SandboxSet { + return newSandboxSet(cfg) +} diff --git a/go/worker/sandboxset/sandboxset.go b/go/worker/sandboxset/sandboxset.go new file mode 100644 index 000000000..f4fd89abe --- /dev/null +++ b/go/worker/sandboxset/sandboxset.go @@ -0,0 +1,182 @@ +package sandboxset + +import ( + "errors" + "fmt" + "log/slog" + "sync" + + "github.com/open-lambda/open-lambda/go/worker/sandbox" +) + +// ErrClosed is returned by operations on a closed SandboxSet. Match with errors.Is. +var ErrClosed = errors.New("sandboxset: closed") + +type SandboxRef struct { + set *sandboxSetImpl + sb sandbox.Sandbox + inUse bool +} + +// Sandbox returns the underlying sandbox. No inUse guard: hot path; misuse is caught by Put/MarkDead guards instead. +func (r *SandboxRef) Sandbox() sandbox.Sandbox { return r.sb } + +func (r *SandboxRef) MarkDead() { + if !r.inUse { + panic(fmt.Sprintf("sandboxset: MarkDead on ref %p not currently held (inUse=%v)", r, r.inUse)) + } + r.sb = nil +} + +func (r *SandboxRef) Put() { r.set.put(r) } + +type sandboxSetImpl struct { + cfg *Config + + mu sync.Mutex + pool []*SandboxRef + closed bool +} + +func newSandboxSet(cfg *Config) *sandboxSetImpl { + if cfg == nil { + panic("sandboxset: Config must not be nil") + } + if cfg.Pool == nil { + panic("sandboxset: Config.Pool must not be nil") + } + if cfg.CodeDir == "" { + panic("sandboxset: Config.CodeDir must not be empty") + } + if cfg.ScratchDirs == nil { + panic("sandboxset: Config.ScratchDirs must not be nil") + } + return &sandboxSetImpl{cfg: cfg} +} + +func (s *sandboxSetImpl) claimIdle() (*SandboxRef, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return nil, fmt.Errorf("claimIdle: %w", ErrClosed) + } + + var empty *SandboxRef + for _, ref := range s.pool { + if ref.inUse { + continue + } + if ref.sb != nil { + ref.inUse = true + return ref, nil + } + if empty == nil { + empty = ref + } + } + + if empty != nil { + empty.inUse = true + return empty, nil + } + + ref := &SandboxRef{set: s, inUse: true} + s.pool = append(s.pool, ref) + return ref, nil +} + +// createSandbox must be called without holding s.mu. +func (s *sandboxSetImpl) createSandbox() (sandbox.Sandbox, error) { + var parentSb sandbox.Sandbox + if s.cfg.Parent != nil { + parentRef, err := s.cfg.Parent.GetOrCreateUnpaused() + if err != nil { + return nil, err + } + parentSb = parentRef.Sandbox() + defer parentRef.Put() + } + + scratchDir := s.cfg.ScratchDirs.Make("sb") + sb, err := s.cfg.Pool.Create( + parentSb, s.cfg.IsLeaf, + s.cfg.CodeDir, scratchDir, + s.cfg.Meta, + ) + if err != nil { + return nil, fmt.Errorf("sandboxset: pool create: %w", err) + } + return sb, nil +} + +func (s *sandboxSetImpl) GetOrCreateUnpaused() (*SandboxRef, error) { + ref, err := s.claimIdle() + if err != nil { + return nil, err + } + + if ref.sb != nil { + if err := ref.sb.Unpause(); err != nil { + slog.Warn("sandboxset: unpause failed, discarding sandbox", "err", err) + ref.sb = nil + } + } + + if ref.sb == nil { + newSb, err := s.createSandbox() + if err != nil { + s.mu.Lock() + ref.inUse = false + s.mu.Unlock() + return nil, err + } + ref.sb = newSb + } + + return ref, nil +} + +func (s *sandboxSetImpl) put(ref *SandboxRef) { + if ref.sb != nil { + if err := ref.sb.Pause(); err != nil { + ref.sb = nil + } + } + + s.mu.Lock() + defer s.mu.Unlock() + + if !ref.inUse { + panic(fmt.Sprintf("sandboxset: put on ref %p not currently held (inUse=%v)", ref, ref.inUse)) + } + if s.closed { + sb := ref.sb + ref.sb = nil + ref.inUse = false + if sb != nil { + sb.Destroy("sandboxset: closed during put") + } + return + } + ref.inUse = false +} + + +func (s *sandboxSetImpl) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return fmt.Errorf("Close: already %w", ErrClosed) + } + s.closed = true + + for _, ref := range s.pool { + if !ref.inUse && ref.sb != nil { + ref.sb.Destroy("sandboxset: closed") + ref.sb = nil + } + } + return nil +} diff --git a/go/worker/sandboxset/tests/sandboxset_integration_test.go b/go/worker/sandboxset/tests/sandboxset_integration_test.go new file mode 100644 index 000000000..a699da323 --- /dev/null +++ b/go/worker/sandboxset/tests/sandboxset_integration_test.go @@ -0,0 +1,189 @@ +//go:build integration + +package tests + +import ( + "os" + "path/filepath" + "testing" + + "github.com/open-lambda/open-lambda/go/common" + "github.com/open-lambda/open-lambda/go/worker/sandbox" + "github.com/open-lambda/open-lambda/go/worker/sandboxset" +) + +// newDockerSet creates a SandboxSet backed by a real DockerPool. +// Requires Docker daemon running and the ol-min image available. +func newDockerSet(t *testing.T) sandboxset.SandboxSet { + t.Helper() + + tmpDir := t.TempDir() + workerDir := filepath.Join(tmpDir, "worker") + pkgsDir := filepath.Join(tmpDir, "packages") + codeDir := filepath.Join(tmpDir, "code") + + for _, d := range []string{workerDir, pkgsDir, codeDir} { + if err := os.MkdirAll(d, 0755); err != nil { + t.Fatal(err) + } + } + + common.Conf = &common.Config{ + Worker_dir: workerDir, + Pkgs_dir: pkgsDir, + Sandbox: "docker", + Docker: common.DockerConfig{ + Base_image: "ol-min", + }, + Limits: common.LimitsConfig{ + Procs: 10, + Mem_mb: 50, + CPU_percent: 100, + Swappiness: 0, + Runtime_sec: 30, + }, + } + + pool, err := sandbox.NewDockerPool("", nil) + if err != nil { + t.Fatalf("NewDockerPool: %v (is Docker running? is ol-min image built?)", err) + } + + scratchDirs, err := common.NewDirMaker("scratch", common.STORE_REGULAR) + if err != nil { + t.Fatal(err) + } + + set := sandboxset.New(&sandboxset.Config{ + Pool: pool, + IsLeaf: true, + CodeDir: codeDir, + ScratchDirs: scratchDirs, + }) + + t.Cleanup(func() { + _ = set.Close() + pool.Cleanup() + }) + + return set +} + +func TestIntegration_GetCreatesRealContainer(t *testing.T) { + set := newDockerSet(t) + + ref, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("GetOrCreateUnpaused: %v", err) + } + sb := ref.Sandbox() + + if sb.ID() == "" { + t.Fatal("expected non-empty sandbox ID") + } + t.Logf("created real container: ID=%s", sb.ID()) + t.Logf("debug: %s", sb.DebugString()) + + // Caller owns lifecycle: destroy the sandbox, mark the ref dead, release. + ref.Sandbox().Destroy("test cleanup") + ref.MarkDead() + ref.Put() +} + +func TestIntegration_PutPausesAndReuses(t *testing.T) { + set := newDockerSet(t) + + // Get a sandbox, record its ID, put it back + ref1, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("first GetOrCreateUnpaused: %v", err) + } + id1 := ref1.Sandbox().ID() + t.Logf("first sandbox: ID=%s", id1) + + ref1.Put() + + // Get again — should reuse the same container (unpaused from paused state) + ref2, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("second GetOrCreateUnpaused: %v", err) + } + id2 := ref2.Sandbox().ID() + t.Logf("second sandbox: ID=%s", id2) + + if id2 != id1 { + t.Fatalf("expected reuse (same ID %s), got new container %s", id1, id2) + } + + ref2.Sandbox().Destroy("test cleanup") + ref2.MarkDead() + ref2.Put() +} + +func TestIntegration_MarkDeadGetsNew(t *testing.T) { + set := newDockerSet(t) + + ref1, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("first GetOrCreateUnpaused: %v", err) + } + id1 := ref1.Sandbox().ID() + t.Logf("first sandbox: ID=%s", id1) + + // Caller-owned destroy + MarkDead + Put releases the slot without a sandbox. + ref1.Sandbox().Destroy("test: simulate handler failure") + ref1.MarkDead() + ref1.Put() + + // Get again — must be a different container since the slot is empty + ref2, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("second GetOrCreateUnpaused: %v", err) + } + id2 := ref2.Sandbox().ID() + t.Logf("second sandbox: ID=%s", id2) + + if id2 == id1 { + t.Fatal("expected new container after MarkDead, got same ID") + } + + ref2.Sandbox().Destroy("test cleanup") + ref2.MarkDead() + ref2.Put() +} + +func TestIntegration_CloseDestroysAll(t *testing.T) { + set := newDockerSet(t) + + // Create multiple sandboxes + refs := make([]*sandboxset.SandboxRef, 3) + for i := range refs { + ref, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("GetOrCreateUnpaused[%d]: %v", i, err) + } + refs[i] = ref + t.Logf("sandbox[%d]: ID=%s", i, ref.Sandbox().ID()) + } + + // Put one back to idle so Close covers the idle path. + refs[2].Put() + + // Close destroys the idle sandbox; in-use refs are destroyed by put() + // when their holders return them below. + if err := set.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + // Put after Close routes through put()'s closed branch, which destroys + // the sandbox. The caller does not Destroy here. + for i := 0; i < 2; i++ { + refs[i].Put() + } + + // Verify set is closed — further Gets should fail + _, err := set.GetOrCreateUnpaused() + if err == nil { + t.Fatal("expected error after Close") + } +} \ No newline at end of file diff --git a/go/worker/sandboxset/tests/sandboxset_test.go b/go/worker/sandboxset/tests/sandboxset_test.go new file mode 100644 index 000000000..647adf8f2 --- /dev/null +++ b/go/worker/sandboxset/tests/sandboxset_test.go @@ -0,0 +1,190 @@ +package tests + +import ( + "testing" + + "github.com/open-lambda/open-lambda/go/common" + "github.com/open-lambda/open-lambda/go/worker/sandbox" + "github.com/open-lambda/open-lambda/go/worker/sandboxset" +) + +// newTestSet creates a valid SandboxSet backed by a MockSandboxPool. +func newTestSet(t *testing.T) (sandboxset.SandboxSet, *sandbox.MockSandboxPool) { + t.Helper() + tmpDir := t.TempDir() + common.Conf = &common.Config{Worker_dir: tmpDir} + scratchDirs, err := common.NewDirMaker("scratch", common.STORE_REGULAR) + if err != nil { + t.Fatal(err) + } + pool := &sandbox.MockSandboxPool{} + set := sandboxset.New(&sandboxset.Config{ + Pool: pool, + CodeDir: tmpDir + "/code", + ScratchDirs: scratchDirs, + }) + return set, pool +} + +// TestGet_CreatesNew verifies that GetOrCreateUnpaused creates a new sandbox +// when the pool is empty. +func TestGet_CreatesNew(t *testing.T) { + set, pool := newTestSet(t) + + ref, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("GetOrCreateUnpaused: %v", err) + } + if ref.Sandbox() == nil { + t.Fatal("expected non-nil sandbox") + } + if n := len(pool.CreatedSandboxes()); n != 1 { + t.Fatalf("expected 1 created sandbox, got %d", n) + } +} + +// TestLifecycle_GetPutReuse verifies the full create → put → reuse cycle. +func TestLifecycle_GetPutReuse(t *testing.T) { + set, _ := newTestSet(t) + defer set.Close() + + ref1, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("GetOrCreateUnpaused: %v", err) + } + id := ref1.Sandbox().ID() + + ref1.Put() + + ref2, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("second GetOrCreateUnpaused: %v", err) + } + if ref2.Sandbox().ID() != id { + t.Fatalf("expected reuse (ID %s), got new (ID %s)", id, ref2.Sandbox().ID()) + } + ref2.Put() +} + +// TestMarkDead_NewSandboxOnNextGet verifies that after MarkDead+Put the slot +// is empty, and the next GetOrCreateUnpaused creates a fresh sandbox in it +// (pool size stays 1, total created sandboxes grows to 2). +func TestMarkDead_NewSandboxOnNextGet(t *testing.T) { + set, pool := newTestSet(t) + defer set.Close() + + ref1, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("GetOrCreateUnpaused: %v", err) + } + id1 := ref1.Sandbox().ID() + + ref1.MarkDead() + ref1.Put() + + ref2, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("second GetOrCreateUnpaused: %v", err) + } + if ref2.Sandbox().ID() == id1 { + t.Fatal("expected a new sandbox after MarkDead, got the same ID") + } + if n := len(pool.CreatedSandboxes()); n != 2 { + t.Fatalf("expected 2 created sandboxes total, got %d", n) + } + ref2.Put() +} + +// TestGet_AfterClose verifies that GetOrCreateUnpaused returns an error after Close. +func TestGet_AfterClose(t *testing.T) { + set, _ := newTestSet(t) + + if err := set.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + _, err := set.GetOrCreateUnpaused() + if err == nil { + t.Fatal("expected error after Close") + } +} + +// TestPut_Twice_Panics verifies the double-Put guard: the second Put on a +// ref that's already been returned must panic rather than silently corrupt. +func TestPut_Twice_Panics(t *testing.T) { + set, _ := newTestSet(t) + defer set.Close() + + ref, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("GetOrCreateUnpaused: %v", err) + } + ref.Put() + + defer func() { + if recover() == nil { + t.Fatal("expected panic on double Put") + } + }() + ref.Put() +} + +// TestClose_DestroysIdleSandbox verifies Close destroys idle sandboxes +// (which the set is the only holder of). +func TestClose_DestroysIdleSandbox(t *testing.T) { + set, pool := newTestSet(t) + + ref, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("GetOrCreateUnpaused: %v", err) + } + sb := pool.CreatedSandboxes()[0] + ref.Put() // ref now idle in the pool + + if err := set.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if !sb.IsDestroyed() { + t.Fatal("expected idle sandbox to be destroyed by Close") + } +} + +// TestPut_AfterClose_DestroysSandbox verifies that a Put arriving after Close +// destroys the returned sandbox (the set is the only remaining holder). +func TestPut_AfterClose_DestroysSandbox(t *testing.T) { + set, pool := newTestSet(t) + + ref, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("GetOrCreateUnpaused: %v", err) + } + sb := pool.CreatedSandboxes()[0] + + if err := set.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + ref.Put() + + if !sb.IsDestroyed() { + t.Fatal("expected sandbox to be destroyed by Put after Close") + } +} + +// TestMarkDead_AfterPut_Panics verifies MarkDead is rejected on a ref that +// is no longer held. +func TestMarkDead_AfterPut_Panics(t *testing.T) { + set, _ := newTestSet(t) + defer set.Close() + + ref, err := set.GetOrCreateUnpaused() + if err != nil { + t.Fatalf("GetOrCreateUnpaused: %v", err) + } + ref.Put() + + defer func() { + if recover() == nil { + t.Fatal("expected panic on MarkDead after Put") + } + }() + ref.MarkDead() +} diff --git a/min-image/Dockerfile b/min-image/Dockerfile index a8c7ec0eb..7a80bdc76 100644 --- a/min-image/Dockerfile +++ b/min-image/Dockerfile @@ -1,11 +1,11 @@ FROM ubuntu:22.04 -RUN apt-get -y --fix-missing update -RUN apt-get -y install wget apt-transport-https curl -RUN apt-get -y install python3 python3-dev python3-pip python-is-python3 -RUN apt-get -y install build-essential libseccomp-dev +RUN apt-get update && apt-get -y install \ + wget apt-transport-https curl \ + python3 python3-dev python3-pip python-is-python3 \ + build-essential libseccomp-dev RUN pip3 install --upgrade pip -RUN pip3 install virtualenv requests tornado==6.1.0 +RUN pip3 install virtualenv requests python-dotenv RUN mkdir /runtimes @@ -16,6 +16,7 @@ RUN cd /tmp/py-runtime && python3 setup.py build_ext --inplace RUN mv /tmp/py-runtime/ol.*.so /runtimes/python/ol.so RUN mv /tmp/py-runtime/server.py /runtimes/python/server.py RUN mv /tmp/py-runtime/server_legacy.py /runtimes/python/server_legacy.py +RUN mv /tmp/py-runtime/server_common.py /runtimes/python/server_common.py RUN rm -rf /tmp/py-runtime # for the Docker container engine diff --git a/min-image/runtimes/python/server.py b/min-image/runtimes/python/server.py index b6b1a4e96..ad3072b8d 100644 --- a/min-image/runtimes/python/server.py +++ b/min-image/runtimes/python/server.py @@ -2,99 +2,26 @@ ''' Python runtime for sock ''' -import os, sys, json, argparse, importlib, traceback, time, fcntl, array, socket, struct +import os +import sys +import socket +import struct +import traceback sys.path.append("/usr/local/lib/python3.10/dist-packages") - -import tornado.ioloop -import tornado.web -import tornado.httpserver -import tornado.wsgi -import tornado.netutil +sys.path.append(os.path.dirname(os.path.abspath(__file__))) import ol +from server_common import web_server_on_sock file_sock_path = "/host/ol.sock" file_sock = None bootstrap_path = None + def web_server(): - print(f"server.py: start web server on fd: {file_sock.fileno()}") - sys.path.append('/handler') - - # TODO: as a safeguard, we should add a mechanism so that the - # import doesn't happen until the cgroup move completes, so that a - # malicious child cannot eat up Zygote resources - import f - - class SockFileHandler(tornado.web.RequestHandler): - # TODO: we should consider how are the different requests used in the context of different applications and functions - # and consider what does the validations should look like for example, should we allow POST requests with no payload etc. - def handle_request(self): - try: - data = self.request.body - try: - event = json.loads(data) if data else None - except: - self.set_status(400) # Bad request if JSON parsing fails - self.write(f'bad request data: "{data}"') - return - - result = f.f(event) if event is not None else f.f({}) - self.write(json.dumps(result)) # Return the result as JSON - except Exception: - self.set_status(500) # Internal server error for unhandled exceptions - self.write(traceback.format_exc()) # Include traceback in response - - - # Define methods for each HTTP method - def get(self): - self.handle_request() - - def post(self): - self.handle_request() - - def put(self): - self.handle_request() - - def delete(self): - self.handle_request() - - def patch(self): - self.handle_request() - - def options(self): - self.handle_request() - - - if hasattr(f, "app"): - def path_wrapper(environ, start_response): - path = environ.get("PATH_INFO", "") - # split path to get individual components - parts = path.split("/") # ["", "run", ] - - # set new environment path - # `/run//a/b/c` -> `/a/b/c` - environ["PATH_INFO"] = '/' + '/'.join(parts[3:]) - - # set the root of the application - app_name = parts[2] - environ["SCRIPT_NAME"] = '/run/' + app_name - - return f.app(environ, start_response) - - # use WSGI entry - # call wrapper to strip /run/ from path - app = tornado.wsgi.WSGIContainer(path_wrapper) - else: - # use function entry - app = tornado.web.Application([ - (".*", SockFileHandler), - ]) - server = tornado.httpserver.HTTPServer(app) - server.add_socket(file_sock) - tornado.ioloop.IOLoop.instance().start() - server.start() + """Wrapper that calls web_server_on_sock with the global file_sock.""" + web_server_on_sock(file_sock, server_name="server.py") def fork_server(): @@ -163,7 +90,9 @@ def start_container(): # child, which will actually use it. This is so that the parent # can know that once the child exits, it is safe to start sending # messages to the sock file. - file_sock = tornado.netutil.bind_unix_socket(file_sock_path) + file_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + file_sock.bind(file_sock_path) + file_sock.listen(1) # backlog=1: we handle one request at a time, no concurrency pid = os.fork() assert pid >= 0 @@ -185,6 +114,7 @@ def start_container(): print("Exception: " + traceback.format_exc()) print("Problematic Python Code:\n" + code) + def main(): ''' caller is expected to do chroot, because we want to use the diff --git a/min-image/runtimes/python/server_common.py b/min-image/runtimes/python/server_common.py new file mode 100644 index 000000000..6e1a0c863 --- /dev/null +++ b/min-image/runtimes/python/server_common.py @@ -0,0 +1,247 @@ +# pylint: disable=line-too-long,invalid-name,broad-except + +''' +Common code shared between server.py and server_legacy.py +''' + +import os +import sys +import json +import asyncio +import http.client +import importlib +import traceback +from enum import Enum +from urllib.parse import urlparse + +from dotenv import load_dotenv + + +class EntryType(Enum): + FUNC = "func" # f(event) -> result + WSGI = "wsgi" # app(environ, start_response) -> iterable + ASGI = "asgi" # await app(scope, receive, send) + + +class RequestParser: + ''' + Parses an HTTP/1.x request from a connection. + + Uses iso-8859-1 (Latin-1) for the request line per RFC 7230 Section 3.2.4: + "Historically, HTTP has allowed field content with text in the ISO-8859-1 + charset." This encoding also provides a safe 1-to-1 byte-to-codepoint mapping, + ensuring any byte sequence decodes without error. + ''' + def __init__(self, conn): + self.rfile = conn.makefile('rb', buffering=65536) + + # Parse request line: "METHOD /path HTTP/1.1\r\n" + line = self.rfile.readline().decode('iso-8859-1').rstrip('\r\n') + self.command, self.path, self.request_version = line.split(None, 2) + + # Parse headers using stdlib (documented API) + self.headers = http.client.parse_headers(self.rfile) + + self.remaining = int(self.headers.get('Content-Length', 0)) + + def read(self, size=-1): + if size < 0: + size = self.remaining + size = min(size, self.remaining) + data = self.rfile.read(size) + self.remaining -= len(data) + return data + + +def handle_func(conn, request, entry_point): + """Handle direct function calls: f(event) -> result""" + try: + body = request.read() + event = json.loads(body) if body else {} + result = entry_point(event) + response_body = json.dumps(result).encode() + status, status_text = 200, "OK" + content_type = "application/json" + except Exception: + response_body = traceback.format_exc().encode() + status, status_text = 500, "Internal Server Error" + content_type = "text/plain" + + conn.sendall(f"HTTP/1.1 {status} {status_text}\r\n".encode()) + conn.sendall(f"Content-Type: {content_type}\r\n".encode()) + conn.sendall(f"Content-Length: {len(response_body)}\r\n".encode()) + conn.sendall(b"Connection: close\r\n") + conn.sendall(b"\r\n") + conn.sendall(response_body) + + +def handle_wsgi(conn, request, entry_point, app_name, path_info, query_string): + """Handle WSGI apps: app(environ, start_response) -> iterable""" + # Host header is required in HTTP/1.1 (RFC 2616 section 14.23) + # Note: we listen on a Unix socket, so port may not be meaningful + host = request.headers['Host'] + if ':' in host: + server_name, server_port = host.split(':', 1) + else: + server_name, server_port = host, "" + + # WSGI 1.0 (PEP 3333): https://peps.python.org/pep-3333/#environ-variables + environ = { + # CGI variables (required) + "REQUEST_METHOD": request.command, + "SCRIPT_NAME": "/run/" + app_name, + "PATH_INFO": path_info, + "QUERY_STRING": query_string, + "SERVER_NAME": server_name, + "SERVER_PORT": server_port, + "SERVER_PROTOCOL": request.request_version, + # wsgi.* variables (required) + "wsgi.version": (1, 0), # PEP 3333 specifies tuple (1, 0) + "wsgi.url_scheme": "http", + "wsgi.input": request, # request.read() handles Content-Length limiting + "wsgi.errors": sys.stderr, + "wsgi.multithread": False, + "wsgi.multiprocess": False, + "wsgi.run_once": False, + } + # HTTP headers -> environ per CGI spec (RFC 3875 section 4.1.18): + # - Convert to uppercase, replace "-" with "_" + # - Prefix with "HTTP_" except Content-Type and Content-Length + for key, value in request.headers.items(): + key = key.upper().replace("-", "_") + if key in ("CONTENT_TYPE", "CONTENT_LENGTH"): + environ[key] = value + else: + environ["HTTP_" + key] = value + + def start_response(status, response_headers, exc_info=None): + conn.sendall(f"HTTP/1.1 {status}\r\n".encode()) + for name, value in response_headers: + conn.sendall(f"{name}: {value}\r\n".encode()) + conn.sendall(b"Connection: close\r\n") + conn.sendall(b"\r\n") + + result = entry_point(environ, start_response) + for chunk in result: + conn.sendall(chunk) + # PEP 3333: if iterable has close(), server must call it for cleanup + if hasattr(result, 'close'): + result.close() + + +def handle_asgi(conn, request, entry_point, app_name, path_info, query_string): + """Handle ASGI apps: await app(scope, receive, send)""" + # TODO: stream body using more_body flag instead of reading all upfront + body = request.read() + + # ASGI 3.0: https://asgi.readthedocs.io/en/latest/specs/www.html#http-connection-scope + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": request.request_version.split("/")[1], # "HTTP/1.1" -> "1.1" + "method": request.command, + "scheme": "http", + "path": path_info, + "query_string": query_string.encode(), + "root_path": "/run/" + app_name, + "headers": [(k.lower().encode(), v.encode()) for k, v in request.headers.items()], + } + + response_started = False + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + async def send(message): + nonlocal response_started + if message["type"] == "http.response.start": + response_started = True + status = message["status"] + conn.sendall(f"HTTP/1.1 {status} OK\r\n".encode()) + for name, value in message.get("headers", []): + conn.sendall(name + b": " + value + b"\r\n") + conn.sendall(b"Connection: close\r\n") + conn.sendall(b"\r\n") + elif message["type"] == "http.response.body": + conn.sendall(message.get("body", b"")) + + try: + asyncio.run(entry_point(scope, receive, send)) + except Exception: + if not response_started: + error = traceback.format_exc().encode() + conn.sendall(b"HTTP/1.1 500 Internal Server Error\r\n") + conn.sendall(b"Content-Type: text/plain\r\n") + conn.sendall(f"Content-Length: {len(error)}\r\n".encode()) + conn.sendall(b"Connection: close\r\n") + conn.sendall(b"\r\n") + conn.sendall(error) + + +def web_server_on_sock(file_sock, server_name="server"): + """ + Main web server loop. Accepts connections and dispatches to appropriate handler. + + Args: + file_sock: The socket to accept connections on + server_name: Name for logging (e.g., "server.py" or "server_legacy.py") + """ + print(f"{server_name}: start web server on fd: {file_sock.fileno()}") + sys.path.append('/handler') + + # Load environment variables from .env file if it exists + env_path = '/handler/.env' + if os.path.exists(env_path): + load_dotenv(env_path) + print(f"{server_name}: loaded environment variables from {env_path}") + + # Import handler module + entry_file = os.environ.get('OL_ENTRY_FILE', 'f.py') + if not entry_file.endswith('.py'): + raise ValueError(f"OL_ENTRY_FILE must end with .py, got: {entry_file}") + module_name = entry_file[:-3] + handler_module = importlib.import_module(module_name) + + # Determine entry point and type + wsgi_entry = os.environ.get('OL_WSGI_ENTRY') + asgi_entry = os.environ.get('OL_ASGI_ENTRY') + if wsgi_entry: + entry_point = getattr(handler_module, wsgi_entry) + entry_type = EntryType.WSGI + elif asgi_entry: + entry_point = getattr(handler_module, asgi_entry) + entry_type = EntryType.ASGI + elif hasattr(handler_module, 'f'): + entry_point = handler_module.f + entry_type = EntryType.FUNC + elif hasattr(handler_module, 'app'): + entry_point = handler_module.app + # Detect ASGI vs WSGI: ASGI apps have async __call__ + if asyncio.iscoroutinefunction(getattr(entry_point, '__call__', None)): + entry_type = EntryType.ASGI + else: + entry_type = EntryType.WSGI + else: + raise ValueError("No entry point found. Define 'f' or 'app' in your module.") + + print(f"{server_name}: entry_type={entry_type.value}") + + while True: + conn, _ = file_sock.accept() + request = RequestParser(conn) + + # Parse path: `/run//a/b/c` -> app_name, `/a/b/c`, query + parsed = urlparse(request.path) + parts = parsed.path.split("/") # ["", "run", , ...] + app_name = parts[2] + path_info = '/' + '/'.join(parts[3:]) + query_string = parsed.query + + if entry_type == EntryType.FUNC: + handle_func(conn, request, entry_point) + elif entry_type == EntryType.WSGI: + handle_wsgi(conn, request, entry_point, app_name, path_info, query_string) + elif entry_type == EntryType.ASGI: + handle_asgi(conn, request, entry_point, app_name, path_info, query_string) + + conn.close() diff --git a/min-image/runtimes/python/server_legacy.py b/min-image/runtimes/python/server_legacy.py index 4b19156a4..e98330ed3 100644 --- a/min-image/runtimes/python/server_legacy.py +++ b/min-image/runtimes/python/server_legacy.py @@ -9,129 +9,63 @@ import os import sys -import json import argparse import importlib -import traceback +import socket -import tornado.ioloop -import tornado.web -import tornado.httpserver -import tornado.netutil -import tornado.wsgi +sys.path.append(os.path.dirname(os.path.abspath(__file__))) + +from dotenv import load_dotenv +from server_common import web_server_on_sock HOST_DIR = '/host' PKGS_DIR = '/packages' HANDLER_DIR = '/handler' +# Load environment variables from .env file if it exists +env_path = f'{HANDLER_DIR}/.env' +if os.path.exists(env_path): + load_dotenv(env_path) + print(f"server_legacy.py: loaded environment variables from {env_path}") + sys.path.append(PKGS_DIR) sys.path.append(HANDLER_DIR) -FS_PATH = os.path.join(HOST_DIR, 'fs.sock') SOCK_PATH = os.path.join(HOST_DIR, 'ol.sock') +FS_PATH = os.path.join(HOST_DIR, 'fs.sock') STDOUT_PATH = os.path.join(HOST_DIR, 'stdout') STDERR_PATH = os.path.join(HOST_DIR, 'stderr') SERVER_PIPE_PATH = os.path.join(HOST_DIR, 'server_pipe') PROCESSES_DEFAULT = 10 -initialized = False parser = argparse.ArgumentParser(description='Listen and serve cache requests or lambda invocations.') parser.add_argument('--cache', action='store_true', default=False, help='Begin as a cache entry.') -# run after forking into sandbox -def init(): - global initialized, f - if initialized: - return - - # assume submitted .py file is /handler/f.py - import f - - initialized = True - -class SockFileHandler(tornado.web.RequestHandler): - def handle_request(self): - try: - data = self.request.body - try: - event = json.loads(data) if data else None - except: - self.set_status(400) # Bad request if JSON parsing fails - self.write(f'bad request data: "{data}"') - return - - result = f.f(event) if event is not None else f.f({}) - self.write(json.dumps(result)) # Return the result as JSON - except Exception: - self.set_status(500) # Internal server error for unhandled exceptions - self.write(traceback.format_exc()) # Include traceback in response - - - # Define methods for each HTTP method - def get(self): - self.handle_request() - - def post(self): - self.handle_request() - - def put(self): - self.handle_request() - - def delete(self): - self.handle_request() - - def patch(self): - self.handle_request() - - def options(self): - self.handle_request() - - -# listen on sock file with Tornado + def lambda_server(): - init() - if hasattr(f, "app"): - def path_wrapper(environ, start_response): - path = environ.get("PATH_INFO", "") - # split path to get individual components - parts = path.split("/") # ["", "run", ] - - # set new environment path - # `/run//a/b/c` -> `/a/b/c` - environ["PATH_INFO"] = '/' + '/'.join(parts[3:]) - - # set the root of the application - app_name = parts[2] - environ["SCRIPT_NAME"] = '/run/' + app_name - - return f.app(environ, start_response) - - # use WSGI entry - # call wrapper to strip /run/ from path - tornado_app = tornado.wsgi.WSGIContainer(path_wrapper) - else: - # use function entry - tornado_app = tornado.web.Application([ - (".*", SockFileHandler), - ]) - server = tornado.httpserver.HTTPServer(tornado_app) - socket = tornado.netutil.bind_unix_socket(SOCK_PATH) - server.add_socket(socket) - # notify worker server that we are ready through stdout - # flush is necessary, and don't put it after tornado start; won't work + """Start the lambda server on a Unix socket.""" + # Create and bind the socket + if os.path.exists(SOCK_PATH): + os.remove(SOCK_PATH) + file_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + file_sock.bind(SOCK_PATH) + file_sock.listen(1) + + # Notify worker server that we are ready with open(SERVER_PIPE_PATH, 'w', encoding='utf-8') as pipe: pipe.write('ready') - tornado.ioloop.IOLoop.instance().start() - server.start(PROCESSES_DEFAULT) -# listen for fds to forkenter + # Run the web server + web_server_on_sock(file_sock, server_name="server_legacy.py") + + def cache_loop(): + """Listen for fds to forkenter (Docker cache mode).""" import ns signal = "cache" r = -1 - count = 0 # only child meant to serve ever escapes the loop while r != 0 or signal == "cache": if r == 0: @@ -165,22 +99,23 @@ def cache_loop(): print('') flush() - count += 1 - print('SERVING HANDLERS') flush() lambda_server() + def flush(): sys.stdout.flush() sys.stderr.flush() + def redirect(): sys.stdout.close() sys.stderr.close() sys.stdout = open(STDOUT_PATH, 'w') sys.stderr = open(STDERR_PATH, 'w') + if __name__ == '__main__': args = parser.parse_args() redirect() diff --git a/python/src/open_lambda.py b/python/src/open_lambda.py index 8853582a2..3184cc3f9 100644 --- a/python/src/open_lambda.py +++ b/python/src/open_lambda.py @@ -11,14 +11,14 @@ def __init__(self, address="localhost:5000"): self._address = address self._session = Session() - def _post(self, path, data=None): + def _post(self, path, data=None, timeout=None): ''' Issues a _post request to the OL worker ''' - return self._session.post(f'http://{self._address}/{path}', pyjson.dumps(data)) + return self._session.post(f'http://{self._address}/{path}', pyjson.dumps(data), timeout=timeout) - def run(self, fn_name, args, json=True): + def run(self, fn_name, args, json=True, timeout=60): ''' Execute a serverless function ''' - resp = self._post(f"run/{fn_name}", args) + resp = self._post(f"run/{fn_name}", args, timeout=timeout) self._check_status_code(resp, "run") if json: diff --git a/scripts/boss_test.py b/scripts/boss_test.py index 0494efa8c..56454e69b 100644 --- a/scripts/boss_test.py +++ b/scripts/boss_test.py @@ -134,7 +134,9 @@ def verify_lambda_config(lambda_name): "HTTP": [{"Method": "POST"}], "Cron": None, "Kafka": None, - } + }, + "Environment": {}, + "ReuseSandbox": True, } assert actual_config == expected_config, ( f"Lambda config mismatch!\nExpected: {expected_config}\nActual: {actual_config}" @@ -200,11 +202,8 @@ def tester(platform): clear_config() launch_boss(platform) - # Step 1: scale up worker - status = json.loads(boss_get("status")) - assert status["state"]["running"] == 0 + # Step 1: scale to 1 worker (boss may auto-launch 1 on some platforms) scale_workers(1) - assert json.loads(boss_get("status"))["state"]["starting"] == 1 wait_for_workers(1) # Step 2: upload and verify lambda diff --git a/scripts/helper/test.py b/scripts/helper/test.py index 36e3c90a4..1245f8f36 100644 --- a/scripts/helper/test.py +++ b/scripts/helper/test.py @@ -14,7 +14,7 @@ TEST_FILTER = [] TEST_BLOCKLIST = [] -WORKER_TYPE = [] +WORKER_TYPE = None RESULTS = OrderedDict({"runs": []}) START_TIME = None @@ -23,6 +23,9 @@ def set_worker_type(new_val): global WORKER_TYPE WORKER_TYPE = new_val +def get_worker_type(): + return WORKER_TYPE + def set_test_filter(new_val): ''' Sets up the filter for all following tests ''' @@ -114,21 +117,18 @@ def _wrapper(*args, **kwargs): worker = WORKER_TYPE() assert worker print("Worker started") - - if worker: - try: - # run test/benchmark - test_t0 = time() - return_val = func(**kwargs) - test_t1 = time() - result["test_seconds"] = test_t1 - test_t0 - result["pass"] = True - except Exception as err: - print(f"Failed to run test: {err}") - result["pass"] = False - result["errors"].append(traceback.format_exc().split("\n")) - - worker.stop() + try: + # run test/benchmark + test_t0 = time() + return_val = func(**kwargs) + test_t1 = time() + result["test_seconds"] = test_t1 - test_t0 + result["pass"] = True + except Exception as err: + print(f"Failed to run test: {err}") + result["pass"] = False + result["errors"].append(traceback.format_exc().split("\n")) + worker.stop() mounts1 = mounts() if len(mounts0) != len(mounts1): diff --git a/scripts/test.py b/scripts/test.py index 58afd5384..bc876a329 100755 --- a/scripts/test.py +++ b/scripts/test.py @@ -14,6 +14,7 @@ import subprocess from time import time +from datetime import datetime from subprocess import call from multiprocessing import Pool @@ -28,7 +29,8 @@ start_tests, check_test_results, set_worker_type, - test + get_worker_type, + test, ) # You can either install the OpenLambda Python bindings @@ -39,7 +41,6 @@ # These will be set by argparse in main() OL_DIR = None -@test def install_examples_to_worker_registry(): """Install all lambda functions from examples directory to worker registry using admin install""" @@ -49,14 +50,12 @@ def install_examples_to_worker_registry(): if not os.path.exists(examples_dir): print(f"Examples directory not found at {examples_dir}") return - # Get all directories in examples + # Get all directories in examples (each directory is a lambda function) example_functions = [] for item in os.listdir(examples_dir): item_path = os.path.join(examples_dir, item) if os.path.isdir(item_path): - # Check if it has f.py (required for lambda functions) - if os.path.exists(os.path.join(item_path, "f.py")): - example_functions.append(item_path) + example_functions.append(item_path) print(f"Found {len(example_functions)} lambda functions in examples directory") # Install each function using admin install command # Find the ol binary - it should be in the project root @@ -76,8 +75,10 @@ def install_examples_to_worker_registry(): print(f"✓ Successfully installed {func_name}") else: print(f"✗ Failed to install {func_name}: {result.stderr}") + raise Exception(f"install failed for {func_name}") except Exception as e: print(f"✗ Error installing {func_name}: {e}") + raise e print("Finished installing example functions") @@ -154,7 +155,7 @@ def stress_one_lambda_task(args): start, seconds = args pos = 0 while time() < start + seconds: - result = open_lambda.run("echo", pos, json=False) + result = open_lambda.run("echo", pos, json=False, timeout=60) assert_eq(result, str(pos)) pos += 1 return pos @@ -307,6 +308,94 @@ def flask_test(): if r.text != "hi\n": raise ValueError(f"r.text should be 'hi\n', not {repr(r.text)}") +@test +def wsgi_post_echo_test(): + """Test that POST body is properly forwarded to WSGI/Flask apps""" + url = 'http://localhost:5000/run/wsgi-post-echo' + + # Test with plain text body + test_body = "hello world" + r = requests.post(url, data=test_body, headers={"Content-Type": "text/plain"}) + check_status_code(r) + if r.text != test_body: + raise ValueError(f"expected '{test_body}', but got '{r.text}'") + + # Test with JSON body + test_json = '{"key": "value"}' + r = requests.post(url, data=test_json, headers={"Content-Type": "application/json"}) + check_status_code(r) + if r.text != test_json: + raise ValueError(f"expected '{test_json}', but got '{r.text}'") + +@test +def flask_entry_test(): + """Test OL_ENTRY_FILE feature with a Flask app using app.py instead of f.py""" + # Test the index route + url = 'http://localhost:5000/run/flask-entry-test' + print("URL", url) + r = requests.get(url) + print("RESPONSE", r) + + if r.status_code != 200: + raise ValueError(f"expected status code 200, but got {r.status_code}") + if r.text != "Hello from app.py!\n": + raise ValueError(f"r.text should be 'Hello from app.py!\\n', not {repr(r.text)}") + + # Test the info route + url_info = 'http://localhost:5000/run/flask-entry-test/info' + print("URL", url_info) + r = requests.get(url_info) + print("RESPONSE", r) + + if r.status_code != 200: + raise ValueError(f"expected status code 200, but got {r.status_code}") + data = r.json() + if data.get("entry_file") != "app.py": + raise ValueError(f"expected entry_file='app.py', got {data}") + +@test +def fastapi_test(): + """Test ASGI support with FastAPI""" + url = 'http://localhost:5000/run/fastapi-test' + print("URL", url) + r = requests.get(url) + print("RESPONSE", r) + + if r.status_code != 200: + raise ValueError(f"expected status code 200, but got {r.status_code}") + + data = r.json() + if data != {"message": "hello world"}: + raise ValueError(f"expected {{'message': 'hello world'}}, but got {data}") + +@test +def wsgi_entry_test(): + """Test OL_WSGI_ENTRY feature with a WSGI entry point not named 'app'""" + # Test the index route + url = 'http://localhost:5000/run/wsgi-entry-test' + print("URL", url) + r = requests.get(url) + print("RESPONSE", r) + + if r.status_code != 200: + raise ValueError(f"expected status code 200, but got {r.status_code}") + if r.text != "Hello from my_wsgi_app!\n": + raise ValueError(f"r.text should be 'Hello from my_wsgi_app!\\n', not {repr(r.text)}") + + # Test the info route + url_info = 'http://localhost:5000/run/wsgi-entry-test/info' + print("URL", url_info) + r = requests.get(url_info) + print("RESPONSE", r) + + if r.status_code != 200: + raise ValueError(f"expected status code 200, but got {r.status_code}") + data = r.json() + if data.get("entry_point") != "my_wsgi_app": + raise ValueError(f"expected entry_point='my_wsgi_app', got {data}") + if data.get("entry_file") != "main.py": + raise ValueError(f"expected entry_file='main.py', got {data}") + @test def test_http_method_restrictions(): url = 'http://localhost:5000/run/lambda-config-test' @@ -336,8 +425,54 @@ def test_http_method_restrictions(): f"for PUT, not {repr(r.text)}" ) +@test +def env_test(): + """Test that environment variables from ol.yaml are properly loaded""" + open_lambda = OpenLambda() + + # Call the env-test function + result = open_lambda.run("env-test", {}) + + # Verify that all configured environment variables are present + expected_vars = { + "MY_ENV_VAR": "Hello from environment", + "DATABASE_URL": "postgresql://user:pass@localhost/db", + "DEBUG_MODE": "true", + "API_KEY": "secret-key-789", + "CUSTOM_PATH": "/usr/local/bin" + } + + # Check that the configured_env_vars match what we expect + if "configured_env_vars" not in result: + raise ValueError(f"configured_env_vars not found in response: {result}") + + configured = result["configured_env_vars"] + + for key, expected_value in expected_vars.items(): + if key not in configured: + raise ValueError(f"Environment variable {key} not found in response") + if configured[key] != expected_value: + raise ValueError( + f"Environment variable {key}={configured[key]} but expected {expected_value}") + + print(f"✓ All {len(expected_vars)} environment variables loaded correctly") + + # Verify DEBUG_MODE enabled all env vars to be returned + if "all_env_vars" not in result: + raise ValueError("DEBUG_MODE=true but all_env_vars not returned") + + return {"env_vars_tested": len(expected_vars)} + def run_tests(): + worker_type = get_worker_type() + worker = worker_type() + assert worker + print(f"Worker started at {datetime.now().strftime('%I:%M%p').lstrip('0').lower()}") + install_examples_to_worker_registry() + print("Examples installed") + worker.stop() + ping_test() # do smoke tests under various configs @@ -358,8 +493,17 @@ def run_tests(): # make sure we can use WSGI apps based on frameworks like Flask flask_test() + wsgi_post_echo_test() + flask_entry_test() + wsgi_entry_test() test_http_method_restrictions() + # test ASGI support with FastAPI + fastapi_test() + + # test environment variables from ol.yaml + env_test() + # make sure code updates get pulled within the cache time with tempfile.TemporaryDirectory() as reg_dir: with TestConfContext(registry=reg_dir, registry_cache_ms=3000): @@ -424,8 +568,6 @@ def main(): else: raise RuntimeError(f"Invalid worker type {args.worker_type}") - install_examples_to_worker_registry() - start_tests() run_tests() diff --git a/wasm-image/Dockerfile b/wasm-image/Dockerfile index b5d071510..a95167662 100644 --- a/wasm-image/Dockerfile +++ b/wasm-image/Dockerfile @@ -1,6 +1,6 @@ FROM ol-min -RUN apt-get -y install clang +RUN apt-get update && apt-get -y install clang # Setup rust environment (prereq for native runtime) RUN curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain nightly-2025-02-15 diff --git a/wasm-worker/Cargo.lock b/wasm-worker/Cargo.lock index 7f57b628a..1eef641bc 100644 --- a/wasm-worker/Cargo.lock +++ b/wasm-worker/Cargo.lock @@ -185,9 +185,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.10.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f61dac84819c6588b558454b194026eb1f09c293b9036ae9b159e74e73ab6cf9" +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" [[package]] name = "cc"