feat: Support configuring proxy mode

CHANGELOG.md

@@ -6,8 +6,11 @@ All notable changes to this project will be documented in this file.
 
 - The Stackable scaler now ensures that a `TrinoCluster` has changed to `ready` more than 5 seconds
   ago before marking it as `ready` ([#68]).
+- Support two different proxy modes: `ProxyAllCalls` (default) and `ProxyFirstCall`.
+  Read the [proxy modes docs](./docs/proxy-modes.md) for details ([#82]).
 
 [#68]: https://github.com/stackabletech/trino-lb/pull/68
+[#82]: https://github.com/stackabletech/trino-lb/pull/82
 
 ## [0.5.0] - 2025-03-14
 

Cargo.toml

@@ -82,7 +82,7 @@ sqlx = { version = "0.8.2", features = [
 strum = { version = "0.27", features = ["derive"] }
 tokio = "1.39"
 tower = "0.5"
-tower-http = { version = "0.6", features = ["compression-full", "tracing"] }
+tower-http = { version = "0.6", features = ["compression-full", "decompression-full", "tracing"] }
 tracing = "0.1"
 tracing-opentelemetry = "0.25"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
@@ -94,4 +94,4 @@ urlencoding = "2.1"
 indicatif = "0.17"
 
 [patch.crates-io]
-prusto = { git = "https://github.com/sbernauer/prusto.git", branch = "feat/derive-name-attribute" }
+prusto = { git = "https://github.com/sbernauer/prusto.git", branch = "feat/derive-name-attribute-2" }

README.md

@@ -15,6 +15,7 @@ A highly available load balancer with support for queueing, routing and auto-sca
   * [Postgres](./docs/persistence/postgres.md)
 * [Scaling](./docs/scaling/index.md)
   * [Stackable](./docs/scaling/stackable.md)
+* [Proxy modes](./docs/proxy-modes.md)
 
 ## Try it out locally
 

deploy/helm/trino-lb/configs/trino-lb-config.yaml

@@ -1,17 +1,18 @@
 trinoLb:
-  externalAddress: https://5.250.179.64:8443
+  externalAddress: https://85.215.178.111:8443
   tls:
     enabled: true
-    certPemFile: /certificates/cert.pem
-    keyPemFile: /certificates/key.pem
+    certPemFile: /certificates/tls.crt
+    keyPemFile: /certificates/tls.key
   persistence:
     redis:
       clusterMode: true
       endpoint: redis://:[email protected]:6379/
     # postgres:
     #   # helm install postgres bitnami/postgresql --version 13.2.18 --set auth.username=trino-lb,auth.password=trino-lb,auth.database=trino_lb
     #   url: postgres://trino-lb:[email protected]/trino_lb
     #   maxConnections: 10
+  proxyMode: ProxyFirstCall # ProxyFirstCall
   refreshQueryCounterInterval: 30s # It is low for testing purpose and to get frequent running queries metrics
   tracing:
     enabled: true
@@ -56,7 +57,7 @@
       - name: trino-m-2
         endpoint: https://trino-m-2-coordinator-default.default.svc.cluster.local:8443
         credentials: *common-credentials
   # oidc:
   #   maxRunningQueries: 3
   #   trinoClusters:
   #     - name: trino-oidc-1
@@ -101,7 +102,7 @@
 routingFallback: m
 
 clusterAutoscaler:
-  reconcileInterval: 1s
+  # reconcileInterval: 1s
   stackable:
     clusters:
       trino-s-1:

deploy/helm/trino-lb/templates/deployment.yaml

@@ -1,27 +1,30 @@
 apiVersion: apps/v1
 kind: Deployment
 metadata:
   name: {{ .Release.Name }}-deployment
   labels:
     app.kubernetes.io/name: trino-lb
     app.kubernetes.io/instance: {{ .Release.Name }}
 spec:
   replicas: {{ .Values.replicas }}
   selector:
     matchLabels:
       app.kubernetes.io/name: trino-lb
       app.kubernetes.io/instance: {{ .Release.Name }}
   template:
     metadata:
       labels:
         app.kubernetes.io/name: trino-lb
         app.kubernetes.io/instance: {{ .Release.Name }}
     spec:
       containers:
         - name: trino-lb
           image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
           imagePullPolicy: {{ .Values.image.pullPolicy }}
           command: ["trino-lb", "--config-file", "/etc/stackable/trino-lb/config/trino-lb-config.yaml"]
+          # env:
+          #   - name: RUST_LOG
+          #     value: debug
           ports:
             - containerPort: 8443
             - containerPort: 9090
@@ -36,7 +39,22 @@
         - name: config
           secret:
             secretName: {{ .Release.Name }}-config
+        # - name: certificates
+        #   secret:
+        #     secretName: {{ .Release.Name }}-certificates
         - name: certificates
-          secret:
-            secretName: {{ .Release.Name }}-certificates
+          ephemeral:
+            volumeClaimTemplate:
+              metadata:
+                annotations:
+                  secrets.stackable.tech/class: tls
+                  secrets.stackable.tech/scope: pod,service=trino-lb
+                  secrets.stackable.tech/format: tls-pem
+              spec:
+                storageClassName: secrets.stackable.tech
+                accessModes:
+                  - ReadWriteOnce
+                resources:
+                  requests:
+                    storage: "1"
       serviceAccountName: {{ .Release.Name }}

deploy/helm/trino-lb/templates/trinos.yaml

@@ -1,10 +1,10 @@
-# 1 small cluster
-{{- range untilStep 1 2 1 }}
+{{- $trinos := list "trino-s-1" "trino-m-1" "trino-m-2" }}
+{{- range $trinos }}
 ---
 apiVersion: trino.stackable.tech/v1alpha1
 kind: TrinoCluster
 metadata:
-  name: trino-s-{{ . }}
+  name: {{ . }}
   namespace: default
 spec:
   image:
@@ -21,46 +21,47 @@
       config.properties:
         tracing.enabled: "true"
         tracing.exporter.endpoint: http://jaeger-collector.default.svc.cluster.local:4317
+        event-listener.config-files: /tmp/http-event-listener.properties
+    podOverrides: &podOverrides
+      spec:
+        containers:
+          - name: trino
+            volumeMounts:
+              - name: http-event-listener-config
+                mountPath: /tmp/http-event-listener.properties
+                subPath: http-event-listener.properties
+        volumes:
+          - name: http-event-listener-config
+            configMap:
+              name: {{ . }}-http-event-listener-config
     roleGroups:
       default:
         replicas: 1
   workers:
     configOverrides: *configOverrides
+    podOverrides: *podOverrides
     roleGroups:
       default:
         replicas: 1
-{{- end }}
-# 2 medium clusters
-{{- range untilStep 1 3 1 }}
 ---
-apiVersion: trino.stackable.tech/v1alpha1
-kind: TrinoCluster
+apiVersion: v1
+kind: ConfigMap
 metadata:
-  name: trino-m-{{ . }}
+  name: {{ . }}-http-event-listener-config
   namespace: default
-spec:
-  image:
-    productVersion: "455"
-  clusterConfig:
-    catalogLabelSelector:
-      matchLabels:
-        trino: trino
-    listenerClass: external-stable
-    authentication:
-      - authenticationClass: trino-users
-  coordinators:
-    configOverrides: &configOverrides
-      config.properties:
-        tracing.enabled: "true"
-        tracing.exporter.endpoint: http://jaeger-collector.default.svc.cluster.local:4317
-    roleGroups:
-      default:
-        replicas: 1
-  workers:
-    configOverrides: *configOverrides
-    roleGroups:
-      default:
-        replicas: 1
+data:
+  http-event-listener.properties: |
+    event-listener.name=http
+
+    http-event-listener.connect-ingest-uri=https://trino-lb.trino-lb.svc.cluster.local:8443/v1/trino-event-listener
+    http-event-listener.connect-retry-count=10
+    # http-event-listener.connect-http-headers="Trino-Cluster:{{ . }}" # IS BROKEN!
+    http-event-listener.http-client.trust-store-path=/stackable/server_tls/truststore.p12
+    http-event-listener.http-client.trust-store-password=changeit
+
+    http-event-listener.log-created=true
+    http-event-listener.log-completed=true
+    http-event-listener.log-split=false
 {{- end }}
 ---
 apiVersion: v1

docs/design.md

@@ -94,6 +94,12 @@ Currently the following autoscalers are implemented:
 
 Read on the [scaling page](./scaling/index.md) for more details.
 
+## 6. Proxy modes
+
+trino-lb can be configured to either proxy all calls to the underlying Trino clusters or only the initial `POST` request and instruct the client to connect directly to the Trino cluster for the subsequent polling requests.
+
+Read on the [proxy modes page](./proxy-modes.md) for more details.
+
 ## Monitoring
 
 trino-lb emits [OpenTelemetry Metrics](https://opentelemetry.io/docs/concepts/signals/metrics/), which (for now) are only exposed as [Prometheus](https://prometheus.io/) metrics on `http://0.0.0.0:9090/metrics`.
@@ -107,6 +113,7 @@ Also, the cluster `trino-s-1` was started on demand, executed the 60 queries and
 ![Grafana screenshot](./assets/grafana-screenshot.png)
 
 ## Tracing
+
 trino-lb emits [OpenTelemetry Traces](https://opentelemetry.io/docs/concepts/signals/traces/) to [OTLP](https://opentelemetry.io/docs/specs/otel/protocol/) endpoints such as [Jaeger](https://www.jaegertracing.io/).
 When proxy-ing requests to Trino we take care of [OpenTelemetry Propagation](https://opentelemetry.io/docs/instrumentation/js/propagation/), so that the Trino spans will show up within the trino-lb spans.
 This enables nice tracing across trino-lb and trino (as seen in the screenshot below)
@@ -118,6 +125,40 @@ This enables nice tracing across trino-lb and trino (as seen in the screenshot b
 This flowchart represents a Trino client submitting a query.
 It might be send to a Trino clusters or queued if all clusters are full.
 
+### General flow
+
+```mermaid
+sequenceDiagram
+    actor client as Trino client
+    participant lb as trino-lb
+    participant trino as Trino
+
+    client ->>+ lb: Submit query using <br/> POST /v1/statement
+    lb ->> lb: Determine if query should be queued<br/>or which Trino cluster it should be routed to.<br/>See below diagram for details
+    lb ->>- client: # Don't strip space
+
+    loop While queued
+        client ->>+ lb: Poll
+        lb ->>- client: # Don't strip space
+    end
+
+    alt Proxy mode: Proxy all calls
+        loop While query running
+            client ->>+ lb: Poll
+            lb -->+ trino: # Don't strip space
+            trino ->>- lb: # Don't strip space
+            lb ->>- client: # Don<>'t strip space
+        end
+    else Proxy mode: Proxy first call
+        loop While query running
+            client ->>+ trino: Poll
+            trino ->>- client: # Don't strip space
+        end
+    end
+```
+
+### Detailed initial `POST /v1/statement`
+
 ```mermaid
 sequenceDiagram
     actor client as Trino client
@@ -141,7 +182,7 @@ sequenceDiagram
 
     lb ->> lb: Determine cluster with the fewest running queries
 
-    alt All active clusters full
+    alt All active clusters reached query limit or all clusters unhealthy
         lb ->>+ persistence: Store queued query
         persistence ->>- lb: # Don't strip space
 
@@ -156,7 +197,7 @@ sequenceDiagram
         lb ->>+ persistence: Store query
         persistence ->>- lb: # Don't strip space
 
-        lb ->> lb: Change nextUri to point to trino-lb
+        lb ->> lb: Change nextUri to point to trino-lb in case all requests should be proxied
     end
 
     lb ->>- client: Response containing nextUri