From c3e40bcb8ea05cdbece42eeb261accff8be4baf8 Mon Sep 17 00:00:00 2001 From: "d.mykhailovskyi" Date: Tue, 20 Oct 2026 15:12:32 +0300 Subject: [PATCH] Add configurable traffic cloning --- config/README.md | 50 ++ config/config.go | 124 +++++ config/config_test.go | 85 +++ .../testdata/bad.traffic_cloning_percent.yml | 21 + config/testdata/traffic_cloning.yml | 36 ++ .../src/content/docs/configuration/metrics.md | 5 +- .../docs/configuration/traffic-cloning.md | 250 +++++++++ metrics.go | 45 +- proxy.go | 16 + proxy_test.go | 440 +++++++++++++++ scripts/generate_traffic_clone_test_data.py | 388 +++++++++++++ traffic_clone.go | 518 ++++++++++++++++++ 12 files changed, 1976 insertions(+), 2 deletions(-) create mode 100644 config/testdata/bad.traffic_cloning_percent.yml create mode 100644 config/testdata/traffic_cloning.yml create mode 100644 docs/src/content/docs/configuration/traffic-cloning.md create mode 100755 scripts/generate_traffic_clone_test_data.py create mode 100644 traffic_clone.go diff --git a/config/README.md b/config/README.md index 45abf23f..f1697d12 100644 --- a/config/README.md +++ b/config/README.md @@ -41,6 +41,9 @@ connection_pool: max_idle_conns: 100 max_idle_conns_per_host: 2 +# Optional traffic cloning configuration +traffic_cloning: [optional] + server: [optional] @@ -157,6 +160,53 @@ params: value: ``` +### +```yml +# Traffic cloning mirrors accepted requests to one or more additional ClickHouse clusters. +# If this section is omitted, traffic cloning is disabled and chproxy behaves as before. +# +# Setting enable=true without targets is invalid. +enable: | default = false [optional] + +# Directory for temporary request body captures. +# If omitted, the operating system default temporary directory is used. +# The directory is created automatically and must be writable by the chproxy process. +temp_dir: | optional + +# Maximum number of request bodies captured and replayed at the same time. +# If omitted or zero, there is no limit. When the limit is reached, clone is skipped. +max_concurrent_captures: | default = 0 [optional] + +# Maximum single captured request body size. +# If omitted or zero, there is no limit. Larger requests are not cloned. +max_capture_size: | optional + +targets: + - name: + + # Optional source filters. If omitted, requests from all sources are eligible. + # All configured source filters must match for the request to be cloned. + from_cluster: | optional + from_user: | optional + from_cluster_user: | optional + + # Clone destination. The cluster must exist in `clusters`. + # The user must exist in the destination cluster users list. + to_cluster: + to_user: + + # Percentage of matching requests to clone. + # Valid range: 0..100. Default: 100. + percent: | default = 100 [optional] + + # Optional timeout for the cloned request. + timeout: | optional + + # Deprecated: accepted for backwards compatibility. + # Request bodies are captured to temporary files instead. + max_queued_chunks: | default = 16 [optional] +``` + ### ```yml # HTTP server configuration diff --git a/config/config.go b/config/config.go index a8f94c4e..51142b95 100644 --- a/config/config.go +++ b/config/config.go @@ -80,6 +80,8 @@ type Config struct { ConnectionPool ConnectionPool `yaml:"connection_pool,omitempty"` + TrafficCloning TrafficCloning `yaml:"traffic_cloning,omitempty"` + // Allow to proxy ping requests AllowPing bool `yaml:"allow_ping,omitempty"` @@ -535,6 +537,128 @@ func (c *Proxy) UnmarshalYAML(unmarshal func(interface{}) error) error { return checkOverflow(c.XXX, "proxy") } +// TrafficCloning describes optional background cloning of accepted requests +// to one or more ClickHouse clusters. +type TrafficCloning struct { + // Enable allows enabling cloning without relying on targets being present. + Enable bool `yaml:"enable,omitempty"` + + // TempDir is a directory for temporary request body captures. + // If omitted, the operating system default temporary directory is used. + TempDir string `yaml:"temp_dir,omitempty"` + + // MaxConcurrentCaptures limits the number of request bodies being captured + // and replayed to clone targets at the same time. + // If omitted or zero, there is no limit. + MaxConcurrentCaptures int `yaml:"max_concurrent_captures,omitempty"` + + // MaxConcurrentClones limits the total number of concurrent clone send + // goroutines across all targets. If omitted or zero, there is no limit. + MaxConcurrentClones int `yaml:"max_concurrent_clones,omitempty"` + + // MaxCaptureSize limits a single captured request body size. + // If omitted or zero, there is no size limit. + MaxCaptureSize ByteSize `yaml:"max_capture_size,omitempty"` + + // Targets contains clone destinations. + Targets []TrafficCloneTarget `yaml:"targets,omitempty"` + + // Catches all undefined fields and must be empty after parsing. + XXX map[string]interface{} `yaml:",inline"` +} + +// UnmarshalYAML implements the yaml.Unmarshaler interface. +func (tc *TrafficCloning) UnmarshalYAML(unmarshal func(interface{}) error) error { + type plain TrafficCloning + if err := unmarshal((*plain)(tc)); err != nil { + return err + } + + if tc.Enable && len(tc.Targets) == 0 { + return fmt.Errorf("`traffic_cloning.targets` must contain at least one target when traffic cloning is enabled") + } + if tc.MaxConcurrentCaptures < 0 { + return fmt.Errorf("`traffic_cloning.max_concurrent_captures` cannot be negative") + } + if tc.MaxConcurrentClones < 0 { + return fmt.Errorf("`traffic_cloning.max_concurrent_clones` cannot be negative") + } + + return checkOverflow(tc.XXX, "traffic_cloning") +} + +func (tc TrafficCloning) IsEnabled() bool { + return tc.Enable || len(tc.Targets) > 0 +} + +// TrafficCloneTarget describes a single clone destination. +type TrafficCloneTarget struct { + // Name is used in logs and metric labels. + Name string `yaml:"name"` + + // FromCluster limits cloning to requests routed to this source cluster. + // If omitted, requests from all clusters are eligible. + FromCluster string `yaml:"from_cluster,omitempty"` + + // FromUser limits cloning to requests from this chproxy user. + // If omitted, requests from all users are eligible. + FromUser string `yaml:"from_user,omitempty"` + + // FromClusterUser limits cloning to requests proxied as this source cluster_user. + // If omitted, requests from all source cluster users are eligible. + FromClusterUser string `yaml:"from_cluster_user,omitempty"` + + // ToCluster is the cluster where cloned requests are sent. + ToCluster string `yaml:"to_cluster"` + + // ToUser is the cluster_user from ToCluster used for cloned requests. + ToUser string `yaml:"to_user"` + + // Percent controls the sampling rate for this target. + // Defaults to 100. + Percent float64 `yaml:"percent,omitempty"` + + // Timeout limits clone request execution. Zero means no extra timeout. + Timeout Duration `yaml:"timeout,omitempty"` + + // Deprecated: accepted for backwards compatibility with the streaming + // implementation. Request bodies are now captured to temporary files. + MaxQueuedChunks int `yaml:"max_queued_chunks,omitempty"` + + // Catches all undefined fields and must be empty after parsing. + XXX map[string]interface{} `yaml:",inline"` +} + +// UnmarshalYAML implements the yaml.Unmarshaler interface. +func (t *TrafficCloneTarget) UnmarshalYAML(unmarshal func(interface{}) error) error { + *t = TrafficCloneTarget{ + Percent: 100, + MaxQueuedChunks: 16, + } + type plain TrafficCloneTarget + if err := unmarshal((*plain)(t)); err != nil { + return err + } + + if len(t.Name) == 0 { + return fmt.Errorf("`traffic_cloning.target.name` cannot be empty") + } + if len(t.ToCluster) == 0 { + return fmt.Errorf("`traffic_cloning.target.to_cluster` cannot be empty for %q", t.Name) + } + if len(t.ToUser) == 0 { + return fmt.Errorf("`traffic_cloning.target.to_user` cannot be empty for %q", t.Name) + } + if t.Percent < 0 || t.Percent > 100 { + return fmt.Errorf("`traffic_cloning.target.percent` must be between 0 and 100 for %q", t.Name) + } + if t.MaxQueuedChunks <= 0 { + return fmt.Errorf("`traffic_cloning.target.max_queued_chunks` must be positive for %q", t.Name) + } + + return checkOverflow(t.XXX, fmt.Sprintf("traffic_cloning.target %q", t.Name)) +} + // Cluster describes CH cluster configuration // The simplest configuration consists of: // diff --git a/config/config_test.go b/config/config_test.go index 76c1e233..f5dafcef 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -331,6 +331,86 @@ func TestLoadConfig(t *testing.T) { MaxErrorReasonSize: ByteSize(1 << 50), }, }, + { + "traffic cloning", + "testdata/traffic_cloning.yml", + Config{ + Server: Server{ + HTTP: HTTP{ + ListenAddr: ":8080", + TimeoutCfg: TimeoutCfg{ + ReadTimeout: Duration(time.Minute), + WriteTimeout: Duration(3 * time.Minute), + IdleTimeout: Duration(10 * time.Minute), + }, + }, + }, + HackMePlease: true, + Clusters: []Cluster{ + { + Name: "cluster", + Scheme: "http", + Nodes: []string{"127.0.0.1:8123"}, + ClusterUsers: []ClusterUser{ + { + Name: "default", + }, + }, + HeartBeat: HeartBeat{ + Interval: Duration(5 * time.Second), + Timeout: Duration(3 * time.Second), + Request: "/ping", + Response: "Ok.\n", + }, + RetryNumber: 0, + }, + { + Name: "preprod", + Scheme: "http", + Nodes: []string{"127.0.0.2:8123"}, + ClusterUsers: []ClusterUser{ + { + Name: "default", + }, + }, + HeartBeat: HeartBeat{ + Interval: Duration(5 * time.Second), + Timeout: Duration(3 * time.Second), + Request: "/ping", + Response: "Ok.\n", + }, + RetryNumber: 0, + }, + }, + Users: []User{ + { + Name: "default", + ToCluster: "cluster", + ToUser: "default", + MaxExecutionTime: Duration(120 * time.Second), + }, + }, + TrafficCloning: TrafficCloning{ + TempDir: "/tmp/chproxy-traffic-clone", + MaxConcurrentCaptures: 4, + MaxCaptureSize: ByteSize(1 << 30), + Targets: []TrafficCloneTarget{ + { + Name: "preprod", + FromCluster: "cluster", + FromUser: "default", + FromClusterUser: "default", + ToCluster: "preprod", + ToUser: "default", + Percent: 25, + Timeout: Duration(30 * time.Second), + MaxQueuedChunks: 32, + }, + }, + }, + MaxErrorReasonSize: ByteSize(1 << 50), + }, + }, } for _, tc := range testCases { @@ -520,6 +600,11 @@ func TestBadConfig(t *testing.T) { "testdata/bad.max_error_reason_size.yml", "cannot parse byte size \"-10B\": it must be positive float followed by optional units. For example, 1.5Gb, 3T", }, + { + "traffic cloning percent", + "testdata/bad.traffic_cloning_percent.yml", + "`traffic_cloning.target.percent` must be between 0 and 100 for \"preprod\"", + }, } for _, tc := range testCases { diff --git a/config/testdata/bad.traffic_cloning_percent.yml b/config/testdata/bad.traffic_cloning_percent.yml new file mode 100644 index 00000000..657d1478 --- /dev/null +++ b/config/testdata/bad.traffic_cloning_percent.yml @@ -0,0 +1,21 @@ +server: + http: + listen_addr: ":8080" + +clusters: + - name: "cluster" + nodes: ["127.0.0.1:8123"] + users: + - name: "default" + +users: + - name: "default" + to_cluster: "cluster" + to_user: "default" + +traffic_cloning: + targets: + - name: "preprod" + to_cluster: "cluster" + to_user: "default" + percent: 101 diff --git a/config/testdata/traffic_cloning.yml b/config/testdata/traffic_cloning.yml new file mode 100644 index 00000000..3dff731a --- /dev/null +++ b/config/testdata/traffic_cloning.yml @@ -0,0 +1,36 @@ +server: + http: + listen_addr: ":8080" + +hack_me_please: true + +clusters: + - name: "cluster" + nodes: ["127.0.0.1:8123"] + users: + - name: "default" + + - name: "preprod" + nodes: ["127.0.0.2:8123"] + users: + - name: "default" + +users: + - name: "default" + to_cluster: "cluster" + to_user: "default" + +traffic_cloning: + temp_dir: "/tmp/chproxy-traffic-clone" + max_concurrent_captures: 4 + max_capture_size: 1GB + targets: + - name: "preprod" + from_cluster: "cluster" + from_user: "default" + from_cluster_user: "default" + to_cluster: "preprod" + to_user: "default" + percent: 25 + timeout: 30s + max_queued_chunks: 32 diff --git a/docs/src/content/docs/configuration/metrics.md b/docs/src/content/docs/configuration/metrics.md index 26f74e0c..acb3dec6 100644 --- a/docs/src/content/docs/configuration/metrics.md +++ b/docs/src/content/docs/configuration/metrics.md @@ -33,9 +33,12 @@ Metrics are exposed in [Prometheus text format](https://prometheus.io/docs/instr | response_body_bytes_total | Counter | The amount of bytes written to response bodies | `user`, `cluster`, `cluster_user`, `replica`, `cluster_node` | | status_codes_total | Counter | Distribution by response status codes | `user`, `cluster`, `cluster_user`, `replica`, `cluster_node`, `code` | | timeout_request_total | Counter | The number of timed out requests | `user`, `cluster`, `cluster_user`, `replica`, `cluster_node` | +| traffic_clone_duration_seconds | Summary | Traffic clone request duration | `user`, `cluster`, `cluster_user`, `target`, `target_cluster`, `target_cluster_user` | +| traffic_clone_in_flight | Gauge | The number of traffic clone requests in flight | `user`, `cluster`, `cluster_user`, `target`, `target_cluster`, `target_cluster_user` | +| traffic_clone_requests_total | Counter | The number of traffic clone requests by result | `user`, `cluster`, `cluster_user`, `target`, `target_cluster`, `target_cluster_user`, `result` | +| traffic_clone_status_codes_total | Counter | Distribution by status codes returned by traffic clone requests | `user`, `cluster`, `cluster_user`, `target`, `target_cluster`, `target_cluster_user`, `code` | | user_queue_overflow_total | Counter | The number of overflows for per-user request queues | `user`, `cluster`, `cluster_user` | An example of [Grafana's](https://grafana.com) dashboard for `chproxy` metrics is available [here](https://github.com/ContentSquare/chproxy/blob/master/chproxy_overview.json). ![dashboard example](https://user-images.githubusercontent.com/2902918/31392734-b2fd4a18-ade2-11e7-84a9-4aaaac4c10d7.png) - diff --git a/docs/src/content/docs/configuration/traffic-cloning.md b/docs/src/content/docs/configuration/traffic-cloning.md new file mode 100644 index 00000000..60dd3d8e --- /dev/null +++ b/docs/src/content/docs/configuration/traffic-cloning.md @@ -0,0 +1,250 @@ +--- +title: Traffic Cloning +sidebar: + order: 8 +--- + +Traffic cloning allows `chproxy` to mirror accepted requests to one or more additional ClickHouse clusters. +It is intended for keeping a pre-production or validation environment close to production without changing the +response returned to the client. + +The feature is optional. If the `traffic_cloning` section is absent, `chproxy` behaves exactly as before and no +traffic is mirrored. + +## Basic configuration + +Add a top-level `traffic_cloning` section: + +```yml +traffic_cloning: + temp_dir: /var/lib/chproxy/traffic-clone + max_concurrent_captures: 4 + max_capture_size: 1GB + targets: + - name: preprod_raw + from_cluster: production_raw + to_cluster: preprod_raw + to_user: default + percent: 10 + timeout: 30s +``` + +The target cluster must already exist in the regular `clusters` section. The target user must exist in the target +cluster's `users` list. + +```yml +clusters: + - name: production_raw + nodes: ["prod-clickhouse-01:8123"] + users: + - name: default + + - name: preprod_raw + nodes: ["preprod-clickhouse-01:8123"] + users: + - name: default +``` + +## Enabling and disabling + +Traffic cloning is enabled when `traffic_cloning.targets` contains at least one target: + +```yml +traffic_cloning: + targets: + - name: preprod + to_cluster: preprod + to_user: default +``` + +It may also be enabled explicitly: + +```yml +traffic_cloning: + enable: true + targets: + - name: preprod + to_cluster: preprod + to_user: default +``` + +To disable traffic cloning, remove the `traffic_cloning` section or set an empty configuration with no targets: + +```yml +# traffic_cloning section omitted +``` + +Setting `enable: true` without targets is invalid and configuration loading will fail. + +## Source filtering + +By default, a clone target receives sampled traffic from all source clusters. Use `from_cluster` to clone only +traffic that was routed to a specific source cluster: + +```yml +traffic_cloning: + targets: + - name: preprod_raw + from_cluster: production_raw + to_cluster: preprod_raw + to_user: default +``` + +Additional optional filters may be used for narrower routing: + +```yml +traffic_cloning: + targets: + - name: preprod_api_raw + from_cluster: production_raw + from_user: api + from_cluster_user: raw_writer + to_cluster: preprod_raw + to_user: raw_writer +``` + +All configured source filters must match for a request to be cloned. + +## Multiple targets + +The same source traffic can be cloned to multiple clusters: + +```yml +traffic_cloning: + targets: + - name: preprod_raw + from_cluster: production_raw + to_cluster: preprod_raw + to_user: default + percent: 25 + + - name: qa_raw + from_cluster: production_raw + to_cluster: qa_raw + to_user: default + percent: 5 +``` + +Each target makes an independent sampling decision and has its own metrics. + +## Sampling percentage + +`percent` controls how much matching traffic is cloned to a target: + +```yml +percent: 10 +``` + +Valid values are from `0` to `100`. If omitted, the default is `100`. + +## Request execution model + +Clone requests are sent asynchronously. The client response is always produced from the primary ClickHouse request, +not from the cloned request. + +If the clone target fails, times out, or returns a non-200 status code, the primary request continues normally. Clone +failures are reported only through logs and metrics. + +For `POST` requests, the request body is captured to a temporary file while the primary request is being read. Clone +requests are sent only after the body has been fully captured. This prevents partial mirrored `INSERT` requests and +avoids buffering the whole request body in memory. + +The temporary file is removed after all selected clone targets finish. The primary request may still be affected by +local temporary-file I/O, but it does not wait for the mirror ClickHouse cluster to read the request body. + +By default, captures are created in the operating system temporary directory. Use `temp_dir` to place them on a +dedicated filesystem: + +```yml +traffic_cloning: + temp_dir: /var/lib/chproxy/traffic-clone +``` + +The directory is created automatically when the configuration is applied. It must be writable by the `chproxy` +process. If it cannot be created or written to, configuration loading fails. + +On startup, `chproxy` removes any leftover temporary capture files from the configured `temp_dir`. This ensures +that a previous crash or unclean shutdown does not leave stale files behind. + +Use `max_concurrent_captures` to cap how many request bodies may be captured and replayed to clone targets at the +same time: + +```yml +traffic_cloning: + max_concurrent_captures: 4 +``` + +If the limit is reached, matching requests are not cloned and the primary request continues normally. The skipped +clone is reported as `traffic_clone_requests_total{result="skipped"}`. + +Use `max_concurrent_clones` to cap how many clone send goroutines may run at the same time across all targets: + +```yml +traffic_cloning: + max_concurrent_clones: 10 +``` + +If the limit is reached when a request is about to be mirrored, that clone is skipped for the affected target and +recorded as `traffic_clone_requests_total{result="skipped"}`. The primary request is not affected. If omitted or +zero, there is no limit. + +Use `max_capture_size` to avoid mirroring unexpectedly large request bodies: + +```yml +traffic_cloning: + max_capture_size: 1GB +``` + +If the request body is larger than this value, it is not cloned and the primary request continues normally. If omitted +or zero, no size limit is applied. + +`max_queued_chunks` is accepted for backwards compatibility with the previous streaming implementation, but it is not +used by the temporary-file capture path. + +## Timeouts + +`timeout` limits the execution time of a cloned request: + +```yml +timeout: 30s +``` + +If omitted or set to zero, no additional clone-specific timeout is applied. + +## Logging + +Traffic cloning uses the existing `chproxy` logger. + +Debug logs are controlled by the top-level `log_debug` option: + +```yml +log_debug: true +``` + +When debug logging is enabled, clone request completion, failures, and body drops are logged. When `log_debug` is +disabled, the module stays quiet unless regular error logging is triggered elsewhere. + +`log_masks` are still applied by the global logger, so sensitive values can be masked in the same way as for the rest +of `chproxy`. + +## Metrics + +Traffic cloning exposes Prometheus metrics with the same namespace and constant labels configured in `server.metrics`. + +| Name | Type | Description | +| --- | --- | --- | +| `traffic_clone_requests_total` | Counter | Clone requests by result: `success`, `error`, `canceled`, or `skipped` | +| `traffic_clone_status_codes_total` | Counter | Status codes returned by clone targets | +| `traffic_clone_duration_seconds` | Summary | Clone request duration | +| `traffic_clone_in_flight` | Gauge | Clone requests currently in flight | + +Common labels: + +- `user` +- `cluster` +- `cluster_user` +- `target` +- `target_cluster` +- `target_cluster_user` + +`traffic_clone_requests_total` also has a `result` label. `traffic_clone_status_codes_total` also has a `code` label. diff --git a/metrics.go b/metrics.go index 0d1eabec..ae6aaf85 100644 --- a/metrics.go +++ b/metrics.go @@ -37,6 +37,10 @@ var ( configSuccessTime prometheus.Gauge badRequest prometheus.Counter retryRequest *prometheus.CounterVec + trafficCloneRequests *prometheus.CounterVec + trafficCloneStatusCodes *prometheus.CounterVec + trafficCloneDuration *prometheus.SummaryVec + trafficCloneInFlight *prometheus.GaugeVec ) func initMetrics(cfg *config.Config) { @@ -308,6 +312,43 @@ func initMetrics(cfg *config.Config) { }, []string{"user", "cluster", "cluster_user", "replica", "cluster_node"}, ) + trafficCloneRequests = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "traffic_clone_requests_total", + Help: "Total number of traffic clone requests by result", + ConstLabels: constLabels, + }, + []string{"user", "cluster", "cluster_user", "target", "target_cluster", "target_cluster_user", "result"}, + ) + trafficCloneStatusCodes = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Name: "traffic_clone_status_codes_total", + Help: "Distribution by status codes returned for traffic clone requests", + ConstLabels: constLabels, + }, + []string{"user", "cluster", "cluster_user", "target", "target_cluster", "target_cluster_user", "code"}, + ) + trafficCloneDuration = prometheus.NewSummaryVec( + prometheus.SummaryOpts{ + Namespace: namespace, + Name: "traffic_clone_duration_seconds", + Help: "Traffic clone request duration", + ConstLabels: constLabels, + Objectives: map[float64]float64{0.5: 1e-1, 0.9: 1e-2, 0.99: 1e-3, 0.999: 1e-4, 1: 1e-5}, + }, + []string{"user", "cluster", "cluster_user", "target", "target_cluster", "target_cluster_user"}, + ) + trafficCloneInFlight = prometheus.NewGaugeVec( + prometheus.GaugeOpts{ + Namespace: namespace, + Name: "traffic_clone_in_flight", + Help: "The number of traffic clone requests in flight", + ConstLabels: constLabels, + }, + []string{"user", "cluster", "cluster_user", "target", "target_cluster", "target_cluster_user"}, + ) } func registerMetrics(cfg *config.Config) { @@ -321,5 +362,7 @@ func registerMetrics(cfg *config.Config) { cacheHit, cacheMiss, cacheSize, cacheItems, cacheSkipped, requestDuration, proxiedResponseDuration, cachedResponseDuration, canceledRequest, timeoutRequest, - configSuccess, configSuccessTime, badRequest, retryRequest) + configSuccess, configSuccessTime, badRequest, retryRequest, + trafficCloneRequests, trafficCloneStatusCodes, trafficCloneDuration, + trafficCloneInFlight) } diff --git a/proxy.go b/proxy.go index 72ec90bb..de4db8b8 100644 --- a/proxy.go +++ b/proxy.go @@ -43,6 +43,7 @@ type reverseProxy struct { users map[string]*user clusters map[string]*cluster caches map[string]*cache.AsyncCache + trafficCloner *trafficCloner hasWildcarded bool maxIdleConns int maxIdleConnsPerHost int @@ -116,6 +117,8 @@ func (rp *reverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) { bytesWritten: responseBodyBytes.With(s.labels), } + rp.getTrafficCloner().cloneRequest(s, req) + req, origParams := s.decorateRequest(req) // wrap body into cachedReadCloser, so we could obtain the original @@ -636,6 +639,11 @@ func (rp *reverseProxy) applyConfig(cfg *config.Config) error { return err } + trafficCloner, err := newTrafficCloner(cfg.TrafficCloning, clusters, rp.rp.Transport) + if err != nil { + return err + } + rp.maxErrorReasonSize = int64(cfg.MaxErrorReasonSize) caches := make(map[string]*cache.AsyncCache, len(cfg.Caches)) @@ -702,6 +710,7 @@ func (rp *reverseProxy) applyConfig(cfg *config.Config) error { rp.lock.Lock() rp.clusters = clusters rp.users = users + rp.trafficCloner = trafficCloner // Swap is needed for deferred closing of old caches. // See the code above where new caches are created. caches, rp.caches = rp.caches, caches @@ -818,6 +827,13 @@ func (rp *reverseProxy) refreshCacheMetrics() { } } +func (rp *reverseProxy) getTrafficCloner() *trafficCloner { + rp.lock.RLock() + defer rp.lock.RUnlock() + + return rp.trafficCloner +} + // find user, cluster and clusterUser // in case of wildcarded user, cluster user is crafted to use original credentials func (rp *reverseProxy) getUser(name string, password string) (found bool, u *user, c *cluster, cu *clusterUser) { diff --git a/proxy_test.go b/proxy_test.go index aa2f149a..b0858df7 100644 --- a/proxy_test.go +++ b/proxy_test.go @@ -9,6 +9,7 @@ import ( "math/rand" "net" "os" + "path/filepath" "regexp" "runtime" "strings" @@ -199,6 +200,445 @@ func TestNewReverseProxy(t *testing.T) { } } +func TestTrafficCloningSendsMirroredRequest(t *testing.T) { + type capturedRequest struct { + query string + user string + password string + body string + } + + cloneRequests := make(chan capturedRequest, 1) + mainServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + fmt.Fprintln(w, okResponse) + })) + defer mainServer.Close() + + cloneServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("query_id") == "" { + fmt.Fprintln(w, okResponse) + return + } + + body, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("cannot read cloned request body: %s", err) + return + } + user, password, _ := r.BasicAuth() + cloneRequests <- capturedRequest{ + query: r.URL.Query().Get("query"), + user: user, + password: password, + body: string(body), + } + fmt.Fprintln(w, okResponse) + })) + defer cloneServer.Close() + + cfg := &config.Config{ + Clusters: []config.Cluster{ + { + Name: "main", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(mainServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "web", + Password: "web-password", + }, + }, + }, + { + Name: "preprod", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(cloneServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "mirror", + Password: "mirror-password", + }, + }, + }, + }, + Users: []config.User{ + { + Name: defaultUsername, + ToCluster: "main", + ToUser: "web", + }, + }, + TrafficCloning: config.TrafficCloning{ + Targets: []config.TrafficCloneTarget{ + { + Name: "preprod", + FromCluster: "main", + FromUser: defaultUsername, + FromClusterUser: "web", + ToCluster: "preprod", + ToUser: "mirror", + Percent: 100, + MaxQueuedChunks: 16, + }, + }, + }, + MaxErrorReasonSize: config.ByteSize(100 << 20), + } + + p, err := newConfiguredProxy(cfg) + if err != nil { + t.Fatalf("unexpected config error: %s", err) + } + + body := strings.NewReader("FORMAT JSONEachRow\n{\"x\":1}\n") + req := httptest.NewRequest("POST", "http://proxy/?query=INSERT%20INTO%20events", body) + resp := makeCustomRequest(p, req) + if resp.StatusCode != http.StatusOK { + t.Fatalf("unexpected main response status: %d", resp.StatusCode) + } + + select { + case got := <-cloneRequests: + if got.query != "INSERT INTO events" { + t.Fatalf("unexpected cloned query: %q", got.query) + } + if got.user != "mirror" || got.password != "mirror-password" { + t.Fatalf("unexpected cloned credentials: %q/%q", got.user, got.password) + } + if got.body != "FORMAT JSONEachRow\n{\"x\":1}\n" { + t.Fatalf("unexpected cloned body: %q", got.body) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for cloned request") + } +} + +func TestTrafficCloningSendsLargeMirroredRequest(t *testing.T) { + const rows = 100000 + row := "{\"id\":123456789,\"value\":\"traffic-clone-regression\"}\n" + largeBody := strings.Repeat(row, rows) + + cloneRequests := make(chan int, 1) + mainServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + fmt.Fprintln(w, okResponse) + })) + defer mainServer.Close() + + cloneServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("query_id") == "" { + fmt.Fprintln(w, okResponse) + return + } + + body, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("cannot read cloned request body: %s", err) + return + } + cloneRequests <- len(body) + fmt.Fprintln(w, okResponse) + })) + defer cloneServer.Close() + + cfg := &config.Config{ + Clusters: []config.Cluster{ + { + Name: "main", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(mainServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "web", + }, + }, + }, + { + Name: "preprod", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(cloneServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "mirror", + }, + }, + }, + }, + Users: []config.User{ + { + Name: defaultUsername, + ToCluster: "main", + ToUser: "web", + }, + }, + TrafficCloning: config.TrafficCloning{ + Targets: []config.TrafficCloneTarget{ + { + Name: "preprod", + FromCluster: "main", + ToCluster: "preprod", + ToUser: "mirror", + Percent: 100, + }, + }, + }, + MaxErrorReasonSize: config.ByteSize(100 << 20), + } + + p, err := newConfiguredProxy(cfg) + if err != nil { + t.Fatalf("unexpected config error: %s", err) + } + + req := httptest.NewRequest("POST", "http://proxy/?query=INSERT%20INTO%20events%20FORMAT%20JSONEachRow", strings.NewReader(largeBody)) + resp := makeCustomRequest(p, req) + if resp.StatusCode != http.StatusOK { + t.Fatalf("unexpected main response status: %d", resp.StatusCode) + } + + select { + case got := <-cloneRequests: + if got != len(largeBody) { + t.Fatalf("unexpected cloned body length: got %d; want %d", got, len(largeBody)) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for cloned large request") + } +} + +func TestTrafficCloningSkipsRequestOverMaxCaptureSize(t *testing.T) { + largeBody := strings.Repeat("0123456789", 1024) + + cloneRequests := make(chan struct{}, 1) + mainServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + fmt.Fprintln(w, okResponse) + })) + defer mainServer.Close() + + cloneServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("query_id") != "" { + cloneRequests <- struct{}{} + } + fmt.Fprintln(w, okResponse) + })) + defer cloneServer.Close() + + cfg := &config.Config{ + Clusters: []config.Cluster{ + { + Name: "main", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(mainServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "web", + }, + }, + }, + { + Name: "preprod", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(cloneServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "mirror", + }, + }, + }, + }, + Users: []config.User{ + { + Name: defaultUsername, + ToCluster: "main", + ToUser: "web", + }, + }, + TrafficCloning: config.TrafficCloning{ + MaxCaptureSize: config.ByteSize(1024), + Targets: []config.TrafficCloneTarget{ + { + Name: "preprod", + FromCluster: "main", + ToCluster: "preprod", + ToUser: "mirror", + Percent: 100, + }, + }, + }, + MaxErrorReasonSize: config.ByteSize(100 << 20), + } + + p, err := newConfiguredProxy(cfg) + if err != nil { + t.Fatalf("unexpected config error: %s", err) + } + + req := httptest.NewRequest("POST", "http://proxy/?query=INSERT%20INTO%20events", strings.NewReader(largeBody)) + resp := makeCustomRequest(p, req) + if resp.StatusCode != http.StatusOK { + t.Fatalf("unexpected main response status: %d", resp.StatusCode) + } + + select { + case <-cloneRequests: + t.Fatal("unexpected cloned request over max_capture_size") + case <-time.After(100 * time.Millisecond): + } +} + +func TestTrafficCloningCreatesTempDir(t *testing.T) { + tempDir := filepath.Join(t.TempDir(), "traffic-clone") + + mainServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + fmt.Fprintln(w, okResponse) + })) + defer mainServer.Close() + + cloneServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + fmt.Fprintln(w, okResponse) + })) + defer cloneServer.Close() + + cfg := &config.Config{ + Clusters: []config.Cluster{ + { + Name: "main", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(mainServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "web", + }, + }, + }, + { + Name: "preprod", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(cloneServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "mirror", + }, + }, + }, + }, + Users: []config.User{ + { + Name: defaultUsername, + ToCluster: "main", + ToUser: "web", + }, + }, + TrafficCloning: config.TrafficCloning{ + TempDir: tempDir, + Targets: []config.TrafficCloneTarget{ + { + Name: "preprod", + FromCluster: "main", + ToCluster: "preprod", + ToUser: "mirror", + Percent: 100, + }, + }, + }, + MaxErrorReasonSize: config.ByteSize(100 << 20), + } + + if _, err := newConfiguredProxy(cfg); err != nil { + t.Fatalf("unexpected config error: %s", err) + } + + stat, err := os.Stat(tempDir) + if err != nil { + t.Fatalf("expected traffic clone temp dir to be created: %s", err) + } + if !stat.IsDir() { + t.Fatalf("expected traffic clone temp path to be a directory") + } +} + +func TestTrafficCloningSkipsNonMatchingSource(t *testing.T) { + cloneRequests := make(chan struct{}, 1) + mainServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + fmt.Fprintln(w, okResponse) + })) + defer mainServer.Close() + + cloneServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("query_id") != "" { + cloneRequests <- struct{}{} + } + fmt.Fprintln(w, okResponse) + })) + defer cloneServer.Close() + + cfg := &config.Config{ + Clusters: []config.Cluster{ + { + Name: "main", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(mainServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "web", + }, + }, + }, + { + Name: "preprod", + Scheme: "http", + Nodes: []string{strings.TrimPrefix(cloneServer.URL, "http://")}, + ClusterUsers: []config.ClusterUser{ + { + Name: "mirror", + }, + }, + }, + }, + Users: []config.User{ + { + Name: defaultUsername, + ToCluster: "main", + ToUser: "web", + }, + }, + TrafficCloning: config.TrafficCloning{ + Targets: []config.TrafficCloneTarget{ + { + Name: "preprod", + FromCluster: "another-main", + ToCluster: "preprod", + ToUser: "mirror", + Percent: 100, + }, + }, + }, + MaxErrorReasonSize: config.ByteSize(100 << 20), + } + + p, err := newConfiguredProxy(cfg) + if err != nil { + t.Fatalf("unexpected config error: %s", err) + } + + req := httptest.NewRequest("GET", "http://proxy/?query=SELECT%201", nil) + resp := makeCustomRequest(p, req) + if resp.StatusCode != http.StatusOK { + t.Fatalf("unexpected main response status: %d", resp.StatusCode) + } + + select { + case <-cloneRequests: + t.Fatal("unexpected cloned request for non-matching source") + case <-time.After(100 * time.Millisecond): + } +} + var badCfg = &config.Config{ Clusters: []config.Cluster{ { diff --git a/scripts/generate_traffic_clone_test_data.py b/scripts/generate_traffic_clone_test_data.py new file mode 100755 index 00000000..b1a9d0c4 --- /dev/null +++ b/scripts/generate_traffic_clone_test_data.py @@ -0,0 +1,388 @@ +#!/usr/bin/env python3 +"""Generate and verify ClickHouse data for chproxy traffic cloning tests.""" + +import argparse +import base64 +import datetime as dt +import json +import os +import random +import re +import sys +import time +import urllib.error +import urllib.parse +import urllib.request + + +COUNTRIES = ("US", "UA", "DE", "FR", "GB", "PL", "ES", "BR") +DEVICES = ("desktop", "mobile", "tablet", "tv") +PAGES = ("/", "/pricing", "/docs", "/signup", "/dashboard", "/reports") +IDENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +class ClickHouseHTTP: + def __init__(self, base_url, user="", password="", timeout=30.0): + self.base_url = normalize_url(base_url) + self.user = user + self.password = password + self.timeout = timeout + + def execute(self, query, body=b"", query_id=None): + params = {"query": query} + if query_id: + params["query_id"] = query_id + url = add_query_params(self.base_url, params) + headers = {"Content-Type": "text/plain; charset=utf-8"} + if self.user: + token = f"{self.user}:{self.password}".encode("utf-8") + headers["Authorization"] = "Basic " + base64.b64encode(token).decode("ascii") + + request = urllib.request.Request(url, data=body, headers=headers, method="POST") + try: + with urllib.request.urlopen(request, timeout=self.timeout) as response: + return response.read().decode("utf-8", errors="replace") + except urllib.error.HTTPError as exc: + detail = exc.read().decode("utf-8", errors="replace") + raise RuntimeError(f"{self.base_url} returned HTTP {exc.code}: {detail}") from exc + except urllib.error.URLError as exc: + raise RuntimeError(f"cannot reach {self.base_url}: {exc.reason}") from exc + + +def normalize_url(url): + if "://" not in url: + url = "http://" + url + parts = urllib.parse.urlsplit(url) + path = parts.path or "/" + return urllib.parse.urlunsplit((parts.scheme, parts.netloc, path, parts.query, parts.fragment)) + + +def add_query_params(url, params): + parts = urllib.parse.urlsplit(url) + current = urllib.parse.parse_qsl(parts.query, keep_blank_values=True) + current.extend(params.items()) + query = urllib.parse.urlencode(current) + return urllib.parse.urlunsplit((parts.scheme, parts.netloc, parts.path, query, parts.fragment)) + + +def checked_identifier(value): + if not IDENT_RE.match(value): + raise argparse.ArgumentTypeError( + f"{value!r} is not a safe ClickHouse identifier; use letters, digits and underscore" + ) + return value + + +def sql_string(value): + return "'" + value.replace("\\", "\\\\").replace("'", "\\'") + "'" + + +def split_env_list(name): + raw = os.getenv(name, "") + return [item.strip() for item in raw.split(",") if item.strip()] + + +def make_batch_id(): + now = dt.datetime.now(dt.timezone.utc).strftime("%Y%m%dT%H%M%SZ") + suffix = random.SystemRandom().randrange(100000, 999999) + return f"tc_{now}_{suffix}" + + +def create_schema(client, database, table, reset=False): + db = checked_identifier(database) + tbl = checked_identifier(table) + client.execute(f"CREATE DATABASE IF NOT EXISTS {db}") + client.execute( + f""" +CREATE TABLE IF NOT EXISTS {db}.{tbl} +( + batch_id String, + row_id UInt64, + event_time DateTime64(3, 'UTC'), + user_id UInt64, + campaign_id UInt32, + site_id UInt32, + country LowCardinality(String), + device LowCardinality(String), + page String, + revenue Decimal(18, 6), + payload String +) +ENGINE = MergeTree +PARTITION BY toYYYYMM(event_time) +ORDER BY (batch_id, row_id) +""".strip() + ) + if reset: + client.execute(f"TRUNCATE TABLE IF EXISTS {db}.{tbl}") + + +def iter_rows(batch_id, start, count, seed, payload_bytes): + rng = random.Random(seed + start) + base = dt.datetime.now(dt.timezone.utc).replace(microsecond=0) - dt.timedelta(minutes=count) + for offset in range(count): + row_id = start + offset + event_time = base + dt.timedelta(seconds=row_id) + pad = "x" * max(0, payload_bytes - 32) + yield { + "batch_id": batch_id, + "row_id": row_id, + "event_time": event_time.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3], + "user_id": rng.randrange(1, 2_000_000), + "campaign_id": rng.randrange(1, 500), + "site_id": rng.randrange(1, 100), + "country": rng.choice(COUNTRIES), + "device": rng.choice(DEVICES), + "page": rng.choice(PAGES), + "revenue": f"{rng.randrange(0, 2_000_000) / 1_000_000:.6f}", + "payload": f"{batch_id}:{row_id}:{pad}", + } + + +def make_json_each_row(rows): + lines = (json.dumps(row, separators=(",", ":")) for row in rows) + return ("\n".join(lines) + "\n").encode("utf-8") + + +def insert_data(client, database, table, batch_id, rows, batch_size, seed, payload_bytes, sleep_ms): + db = checked_identifier(database) + tbl = checked_identifier(table) + inserted = 0 + while inserted < rows: + size = min(batch_size, rows - inserted) + body = make_json_each_row(iter_rows(batch_id, inserted, size, seed, payload_bytes)) + query_id = f"{batch_id}_batch_{inserted // batch_size:05d}" + client.execute(f"INSERT INTO {db}.{tbl} FORMAT JSONEachRow", body=body, query_id=query_id) + inserted += size + print(f"inserted {inserted}/{rows} rows via chproxy", flush=True) + if sleep_ms: + time.sleep(sleep_ms / 1000.0) + + +def read_stats(client, database, table, batch_id): + db = checked_identifier(database) + tbl = checked_identifier(table) + query = f""" +SELECT + count() AS rows, + uniqExact(row_id) AS unique_rows, + if(count() = 0, 0, min(row_id)) AS min_row_id, + if(count() = 0, 0, max(row_id)) AS max_row_id, + sum(row_id) AS row_id_sum, + round(sum(toFloat64(revenue)), 6) AS revenue_sum +FROM {db}.{tbl} +WHERE batch_id = {sql_string(batch_id)} +FORMAT JSON +""".strip() + raw = client.execute(query) + parsed = json.loads(raw) + data = parsed.get("data") or [{}] + return normalize_stats(data[0]) + + +def normalize_stats(stats): + normalized = {} + for key, value in stats.items(): + if isinstance(value, str) and value.isdigit(): + normalized[key] = int(value) + else: + normalized[key] = value + return normalized + + +def fetch_metrics(metrics_url, timeout): + url = normalize_url(metrics_url) + try: + with urllib.request.urlopen(url, timeout=timeout) as response: + text = response.read().decode("utf-8", errors="replace") + except urllib.error.URLError as exc: + raise RuntimeError(f"cannot fetch metrics from {url}: {exc.reason}") from exc + return [line for line in text.splitlines() if "traffic_clone_" in line and not line.startswith("#")] + + +def endpoint_label(kind, url, index=None): + suffix = "" if index is None else f" #{index}" + return f"{kind}{suffix} ({normalize_url(url)})" + + +def positive_int(value): + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be positive") + return parsed + + +def non_negative_int(value): + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("must be non-negative") + return parsed + + +def percent(value): + parsed = float(value) + if parsed < 0 or parsed > 100: + raise argparse.ArgumentTypeError("must be between 0 and 100") + return parsed + + +def build_parser(): + parser = argparse.ArgumentParser( + description=( + "Create ClickHouse test tables, insert marked rows through chproxy, " + "and verify that traffic-cloned targets received the same batch." + ) + ) + parser.add_argument("--chproxy-url", default=os.getenv("CHPROXY_URL", "http://127.0.0.1:9001/")) + parser.add_argument("--chproxy-user", default=os.getenv("CHPROXY_USER", "default")) + parser.add_argument("--chproxy-password", default=os.getenv("CHPROXY_PASSWORD", "password")) + parser.add_argument( + "--primary-url", + default=os.getenv("CLICKHOUSE_PRIMARY_URL"), + help="optional direct ClickHouse HTTP URL for verifying the primary/source cluster", + ) + parser.add_argument( + "--clone-url", + action="append", + default=None, + help="direct ClickHouse HTTP URL for a clone target; may be passed multiple times", + ) + parser.add_argument("--clickhouse-user", default=os.getenv("CLICKHOUSE_USER", "default")) + parser.add_argument("--clickhouse-password", default=os.getenv("CLICKHOUSE_PASSWORD", "")) + parser.add_argument("--database", type=checked_identifier, default=os.getenv("TC_DATABASE", "traffic_clone_test")) + parser.add_argument("--table", type=checked_identifier, default=os.getenv("TC_TABLE", "events")) + parser.add_argument("--rows", type=positive_int, default=int(os.getenv("TC_ROWS", "10000"))) + parser.add_argument("--batch-size", type=positive_int, default=int(os.getenv("TC_BATCH_SIZE", "1000"))) + parser.add_argument("--batch-id", default=os.getenv("TC_BATCH_ID") or make_batch_id()) + parser.add_argument("--seed", type=int, default=int(os.getenv("TC_SEED", "42"))) + parser.add_argument("--payload-bytes", type=non_negative_int, default=int(os.getenv("TC_PAYLOAD_BYTES", "64"))) + parser.add_argument("--insert-sleep-ms", type=non_negative_int, default=int(os.getenv("TC_INSERT_SLEEP_MS", "0"))) + parser.add_argument("--clone-wait-seconds", type=float, default=float(os.getenv("TC_CLONE_WAIT_SECONDS", "2"))) + parser.add_argument("--timeout", type=float, default=float(os.getenv("TC_HTTP_TIMEOUT", "30"))) + parser.add_argument("--reset", action="store_true", help="truncate the test table on all endpoints before inserting") + parser.add_argument( + "--setup-via-chproxy", + action="store_true", + help=( + "create/truncate the test table through chproxy, so chproxy credentials and " + "traffic cloning create it in primary and clone targets" + ), + ) + parser.add_argument("--skip-setup", action="store_true") + parser.add_argument("--skip-insert", action="store_true") + parser.add_argument("--skip-verify", action="store_true") + parser.add_argument( + "--verify-primary-via-chproxy", + action="store_true", + help="verify primary row count through chproxy when no direct primary ClickHouse URL is available", + ) + parser.add_argument( + "--expect-clone-percent", + type=percent, + default=float(os.getenv("TC_EXPECT_CLONE_PERCENT", "100")), + help="expected clone percentage; sampling is per INSERT request, not per row", + ) + parser.add_argument("--clone-tolerance-rows", type=non_negative_int, default=0) + parser.add_argument("--skip-clone-count-check", action="store_true") + parser.add_argument("--non-strict", action="store_true", help="print verification failures but exit with code 0") + parser.add_argument("--show-metrics", action="store_true") + parser.add_argument("--metrics-url", default=os.getenv("CHPROXY_METRICS_URL")) + return parser + + +def main(): + args = build_parser().parse_args() + clone_urls = args.clone_url if args.clone_url is not None else split_env_list("CLICKHOUSE_CLONE_URLS") + direct_endpoints = [] + if args.primary_url: + direct_endpoints.append(("primary", args.primary_url)) + direct_endpoints.extend((f"clone #{idx}", url) for idx, url in enumerate(clone_urls, 1)) + + chproxy = ClickHouseHTTP(args.chproxy_url, args.chproxy_user, args.chproxy_password, args.timeout) + direct_clients = [ + (label, ClickHouseHTTP(url, args.clickhouse_user, args.clickhouse_password, args.timeout)) + for label, url in direct_endpoints + ] + if args.verify_primary_via_chproxy and not args.primary_url: + direct_clients.insert(0, ("primary via chproxy", chproxy)) + + print(f"batch_id: {args.batch_id}") + print(f"table: {args.database}.{args.table}") + print(f"chproxy: {normalize_url(args.chproxy_url)}") + if args.primary_url: + print(endpoint_label("primary", args.primary_url)) + elif args.verify_primary_via_chproxy: + print("primary: will be verified through chproxy") + else: + print("primary: no --primary-url provided; primary direct verification will be skipped") + if clone_urls: + for idx, url in enumerate(clone_urls, 1): + print(endpoint_label("clone", url, idx)) + else: + print("clone: no --clone-url provided; only primary will be verified") + + if not args.skip_setup: + if args.setup_via_chproxy: + print(f"setting up through chproxy: {chproxy.base_url}") + create_schema(chproxy, args.database, args.table, reset=args.reset) + else: + if not direct_clients: + raise RuntimeError("setup requires --primary-url/--clone-url or --setup-via-chproxy") + for label, client in direct_clients: + if label == "primary via chproxy": + continue + print(f"setting up {label}: {client.base_url}") + create_schema(client, args.database, args.table, reset=args.reset) + + if not args.skip_insert: + insert_data( + chproxy, + args.database, + args.table, + args.batch_id, + args.rows, + args.batch_size, + args.seed, + args.payload_bytes, + args.insert_sleep_ms, + ) + + if args.clone_wait_seconds > 0 and clone_urls and not args.skip_insert: + print(f"waiting {args.clone_wait_seconds:g}s for async clone requests") + time.sleep(args.clone_wait_seconds) + + failures = [] + if not args.skip_verify: + for label, client in direct_clients: + stats = read_stats(client, args.database, args.table, args.batch_id) + rows = int(stats.get("rows", 0)) + print(f"{label}: {json.dumps(stats, sort_keys=True)}") + if label.startswith("primary"): + if rows != args.rows: + failures.append(f"primary has {rows} rows, expected {args.rows}") + elif not args.skip_clone_count_check: + expected = int(round(args.rows * args.expect_clone_percent / 100.0)) + delta = abs(rows - expected) + if delta > args.clone_tolerance_rows: + failures.append( + f"{label} has {rows} rows, expected {expected} +/- {args.clone_tolerance_rows}" + ) + + if args.show_metrics: + metrics_url = args.metrics_url or urllib.parse.urljoin(normalize_url(args.chproxy_url), "metrics") + print(f"traffic clone metrics from {metrics_url}:") + for line in fetch_metrics(metrics_url, args.timeout): + print(line) + + if failures: + print("verification failed:", file=sys.stderr) + for failure in failures: + print(f" - {failure}", file=sys.stderr) + return 0 if args.non_strict else 1 + + print("done") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/traffic_clone.go b/traffic_clone.go new file mode 100644 index 00000000..090ed8bf --- /dev/null +++ b/traffic_clone.go @@ -0,0 +1,518 @@ +package main + +import ( + "context" + "errors" + "fmt" + "io" + "math/rand/v2" + "net/http" + "os" + "path/filepath" + "strconv" + "sync" + "sync/atomic" + "time" + + "github.com/contentsquare/chproxy/config" + "github.com/contentsquare/chproxy/log" + "github.com/prometheus/client_golang/prometheus" +) + +const trafficCloneResponseDrainLimit = 64 << 10 + +type trafficCloner struct { + client *http.Client + tempDir string + maxCaptureSize int64 + captureTokens chan struct{} + cloneTokens chan struct{} + targets []*trafficCloneTarget +} + +type trafficCloneTarget struct { + name string + fromCluster string + fromUser string + fromClusterUser string + cluster *cluster + clusterUser *clusterUser + percent float64 + timeout time.Duration +} + +// cloneSource holds a snapshot of the originating request's identity. +// It is created once from a *scope before goroutines are launched so that +// goroutines do not hold a reference to the scope past the primary request +// lifecycle (which could produce misleading in-flight counters on s.dec()). +type cloneSource struct { + user *user + userName string + clusterName string + cuName string +} + +func newCloneSource(s *scope) cloneSource { + return cloneSource{ + user: s.user, + userName: s.user.name, + clusterName: s.cluster.name, + cuName: s.clusterUser.name, + } +} + +func newTrafficCloner(cfg config.TrafficCloning, clusters map[string]*cluster, transport http.RoundTripper) (*trafficCloner, error) { + if !cfg.IsEnabled() { + return nil, nil + } + if len(cfg.Targets) == 0 { + return nil, fmt.Errorf("`traffic_cloning.targets` must contain at least one target when traffic cloning is enabled") + } + if err := prepareTrafficCloneTempDir(cfg.TempDir); err != nil { + return nil, err + } + + targets := make([]*trafficCloneTarget, 0, len(cfg.Targets)) + for _, tc := range cfg.Targets { + if tc.Name == "" { + return nil, fmt.Errorf("`traffic_cloning.target.name` cannot be empty") + } + c := clusters[tc.ToCluster] + if c == nil { + return nil, fmt.Errorf("unknown `traffic_cloning.target.to_cluster` %q", tc.ToCluster) + } + + cu := c.users[tc.ToUser] + if cu == nil { + return nil, fmt.Errorf("unknown `traffic_cloning.target.to_user` %q in cluster %q", tc.ToUser, tc.ToCluster) + } + if tc.Percent < 0 || tc.Percent > 100 { + return nil, fmt.Errorf("`traffic_cloning.target.percent` must be between 0 and 100 for %q", tc.Name) + } + + targets = append(targets, &trafficCloneTarget{ + name: tc.Name, + fromCluster: tc.FromCluster, + fromUser: tc.FromUser, + fromClusterUser: tc.FromClusterUser, + cluster: c, + clusterUser: cu, + percent: tc.Percent, + timeout: time.Duration(tc.Timeout), + }) + } + + var captureTokens chan struct{} + if cfg.MaxConcurrentCaptures > 0 { + captureTokens = make(chan struct{}, cfg.MaxConcurrentCaptures) + } + + var cloneTokens chan struct{} + if cfg.MaxConcurrentClones > 0 { + cloneTokens = make(chan struct{}, cfg.MaxConcurrentClones) + } + + return &trafficCloner{ + client: &http.Client{ + Transport: transport, + }, + tempDir: cfg.TempDir, + maxCaptureSize: int64(cfg.MaxCaptureSize), + captureTokens: captureTokens, + cloneTokens: cloneTokens, + targets: targets, + }, nil +} + +func prepareTrafficCloneTempDir(dir string) error { + if dir == "" { + return nil + } + if err := os.MkdirAll(dir, 0o750); err != nil { + return fmt.Errorf("cannot create `traffic_cloning.temp_dir` %q: %w", dir, err) + } + + // Remove stale captures left by a previous crash or unclean shutdown. + stale, _ := filepath.Glob(filepath.Join(dir, "chproxy-traffic-clone-*")) + for _, f := range stale { + _ = os.Remove(f) + } + + f, err := os.CreateTemp(dir, "chproxy-traffic-clone-probe-*") + if err != nil { + return fmt.Errorf("cannot write to `traffic_cloning.temp_dir` %q: %w", dir, err) + } + name := f.Name() + if err := f.Close(); err != nil { + _ = os.Remove(name) + return fmt.Errorf("cannot close probe file in `traffic_cloning.temp_dir` %q: %w", dir, err) + } + if err := os.Remove(name); err != nil { + return fmt.Errorf("cannot remove probe file from `traffic_cloning.temp_dir` %q: %w", dir, err) + } + return nil +} + +func (tc *trafficCloner) cloneRequest(source *scope, req *http.Request) { + if tc == nil { + return + } + + targets := tc.pickTargets(source) + if len(targets) == 0 { + return + } + + if req.Body != nil && req.Body != http.NoBody { + if tc.maxCaptureSize > 0 && req.ContentLength > tc.maxCaptureSize { + tc.skipTargets(source, targets, "request body is larger than traffic_cloning.max_capture_size") + return + } + if !tc.tryAcquireCapture() { + tc.skipTargets(source, targets, "traffic clone capture concurrency limit reached") + return + } + capture, err := newTrafficCloneBodyCapture(req.Body, tc, source, targets, req) + if err != nil { + tc.releaseCapture() + src := newCloneSource(source) + for _, target := range targets { + labels := trafficCloneLabels(src, target) + trafficCloneRequests.With(trafficCloneResultLabels(labels, "error")).Inc() + log.Errorf("%s: cannot initialize traffic clone capture for target %q: %s", src.userName, target.name, err) + } + return + } + req.Body = capture + return + } + + tc.cloneRequestsWithoutBody(source, targets, req) +} + +func (tc *trafficCloner) tryAcquireCapture() bool { + if tc.captureTokens == nil { + return true + } + select { + case tc.captureTokens <- struct{}{}: + return true + default: + return false + } +} + +func (tc *trafficCloner) releaseCapture() { + if tc.captureTokens == nil { + return + } + <-tc.captureTokens +} + +func (tc *trafficCloner) tryAcquireClone() bool { + if tc.cloneTokens == nil { + return true + } + select { + case tc.cloneTokens <- struct{}{}: + return true + default: + return false + } +} + +func (tc *trafficCloner) releaseClone() { + if tc.cloneTokens == nil { + return + } + <-tc.cloneTokens +} + +func (tc *trafficCloner) skipTargets(source *scope, targets []*trafficCloneTarget, reason string) { + src := newCloneSource(source) + for _, target := range targets { + labels := trafficCloneLabels(src, target) + trafficCloneRequests.With(trafficCloneResultLabels(labels, "skipped")).Inc() + log.Debugf("%s: traffic clone to target %q skipped: %s", src.userName, target.name, reason) + } +} + +func (tc *trafficCloner) pickTargets(source *scope) []*trafficCloneTarget { + targets := make([]*trafficCloneTarget, 0, len(tc.targets)) + for _, target := range tc.targets { + if target.matches(source) && (target.percent >= 100 || tc.randomPercent() < target.percent) { + targets = append(targets, target) + } + } + return targets +} + +func (target *trafficCloneTarget) matches(source *scope) bool { + if target.fromCluster != "" && target.fromCluster != source.cluster.name { + return false + } + if target.fromUser != "" && target.fromUser != source.user.name { + return false + } + if target.fromClusterUser != "" && target.fromClusterUser != source.clusterUser.name { + return false + } + return true +} + +func (tc *trafficCloner) randomPercent() float64 { + return rand.Float64() * 100 +} + +func (tc *trafficCloner) cloneRequestsWithoutBody(source *scope, targets []*trafficCloneTarget, req *http.Request) { + src := newCloneSource(source) + for _, target := range targets { + if !tc.tryAcquireClone() { + labels := trafficCloneLabels(src, target) + trafficCloneRequests.With(trafficCloneResultLabels(labels, "skipped")).Inc() + log.Debugf("%s: traffic clone to target %q skipped: clone concurrency limit reached", src.userName, target.name) + continue + } + cloneReq := req.Clone(context.Background()) + cloneReq.RequestURI = "" + go func() { + defer tc.releaseClone() + tc.send(src, target, cloneReq) + }() + } +} + +func (tc *trafficCloner) cloneCapturedRequests(source cloneSource, targets []*trafficCloneTarget, requests []*http.Request, fileName string) { + defer tc.releaseCapture() + + var wg sync.WaitGroup + for i, target := range targets { + target := target + cloneReq := requests[i] + + if !tc.tryAcquireClone() { + labels := trafficCloneLabels(source, target) + trafficCloneRequests.With(trafficCloneResultLabels(labels, "skipped")).Inc() + log.Debugf("%s: traffic clone to target %q skipped: clone concurrency limit reached", source.userName, target.name) + continue + } + + wg.Add(1) + go func() { + defer wg.Done() + defer tc.releaseClone() + + f, err := os.Open(fileName) + if err != nil { + labels := trafficCloneLabels(source, target) + trafficCloneRequests.With(trafficCloneResultLabels(labels, "error")).Inc() + log.Errorf("%s: cannot open traffic clone capture for target %q: %s", source.userName, target.name, err) + return + } + defer f.Close() + cloneReq.Body = f + tc.send(source, target, cloneReq) + }() + } + wg.Wait() + if err := os.Remove(fileName); err != nil { + log.Errorf("cannot remove traffic clone capture %q: %s", fileName, err) + } +} + +func (tc *trafficCloner) send(source cloneSource, target *trafficCloneTarget, req *http.Request) { + ctx := context.Background() + var cancel context.CancelFunc + if target.timeout > 0 { + ctx, cancel = context.WithTimeout(ctx, target.timeout) + } else { + ctx, cancel = context.WithCancel(ctx) + } + defer cancel() + + req = req.WithContext(ctx) + cloneScope := newScope(req, source.user, target.cluster, target.clusterUser, "", getSessionTimeout(req)) + cloneScope.decorateRequest(req) + + labels := trafficCloneLabels(source, target) + trafficCloneInFlight.With(labels).Inc() + defer trafficCloneInFlight.With(labels).Dec() + + startTime := time.Now() + resp, err := tc.client.Do(req) + trafficCloneDuration.With(labels).Observe(time.Since(startTime).Seconds()) + if err != nil { + result := "error" + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + result = "canceled" + } + trafficCloneRequests.With(trafficCloneResultLabels(labels, result)).Inc() + log.Debugf("%s: traffic clone to target %q failed: %s", source.userName, target.name, err) + return + } + defer resp.Body.Close() + + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, trafficCloneResponseDrainLimit)) + + trafficCloneRequests.With(trafficCloneResultLabels(labels, "success")).Inc() + trafficCloneStatusCodes.With(trafficCloneStatusCodeLabels(labels, resp.StatusCode)).Inc() + log.Debugf("%s: traffic clone to target %q completed with status code %d", source.userName, target.name, resp.StatusCode) +} + +func trafficCloneLabels(source cloneSource, target *trafficCloneTarget) prometheus.Labels { + return prometheus.Labels{ + "user": source.userName, + "cluster": source.clusterName, + "cluster_user": source.cuName, + "target": target.name, + "target_cluster": target.cluster.name, + "target_cluster_user": target.clusterUser.name, + } +} + +func trafficCloneResultLabels(labels prometheus.Labels, result string) prometheus.Labels { + withResult := prometheus.Labels{} + for k, v := range labels { + withResult[k] = v + } + withResult["result"] = result + return withResult +} + +func trafficCloneStatusCodeLabels(labels prometheus.Labels, code int) prometheus.Labels { + withCode := prometheus.Labels{} + for k, v := range labels { + withCode[k] = v + } + withCode["code"] = strconv.Itoa(code) + return withCode +} + +type trafficCloneBodyCapture struct { + io.ReadCloser + + cloner *trafficCloner + source cloneSource + targets []*trafficCloneTarget + + file *os.File + fileName string + requests []*http.Request + captured int64 + + done sync.Once + finalized atomic.Bool +} + +func newTrafficCloneBodyCapture( + body io.ReadCloser, + cloner *trafficCloner, + source *scope, + targets []*trafficCloneTarget, + req *http.Request, +) (*trafficCloneBodyCapture, error) { + f, err := os.CreateTemp(cloner.tempDir, "chproxy-traffic-clone-*") + if err != nil { + return nil, err + } + + requests := make([]*http.Request, len(targets)) + for i := range targets { + cloneReq := req.Clone(context.Background()) + cloneReq.RequestURI = "" + cloneReq.Body = nil + requests[i] = cloneReq + } + + return &trafficCloneBodyCapture{ + ReadCloser: body, + cloner: cloner, + source: newCloneSource(source), + targets: targets, + file: f, + fileName: f.Name(), + requests: requests, + }, nil +} + +func (c *trafficCloneBodyCapture) Read(p []byte) (int, error) { + n, err := c.ReadCloser.Read(p) + if c.finalized.Load() { + return n, err + } + if n > 0 { + c.captured += int64(n) + if c.cloner.maxCaptureSize > 0 && c.captured > c.cloner.maxCaptureSize { + c.skip("request body exceeded traffic_cloning.max_capture_size") + return n, err + } + if _, werr := c.file.Write(p[:n]); werr != nil { + c.fail(werr) + } + } + if err == io.EOF { + c.finish() + } else if err != nil { + c.abort() + } + return n, err +} + +func (c *trafficCloneBodyCapture) Close() error { + c.abort() + return c.ReadCloser.Close() +} + +func (c *trafficCloneBodyCapture) finish() { + c.done.Do(func() { + c.finalized.Store(true) + if err := c.file.Close(); err != nil { + _ = os.Remove(c.fileName) + for _, target := range c.targets { + labels := trafficCloneLabels(c.source, target) + trafficCloneRequests.With(trafficCloneResultLabels(labels, "error")).Inc() + log.Errorf("%s: cannot close traffic clone capture for target %q: %s", c.source.userName, target.name, err) + } + c.cloner.releaseCapture() + return + } + go c.cloner.cloneCapturedRequests(c.source, c.targets, c.requests, c.fileName) + }) +} + +func (c *trafficCloneBodyCapture) abort() { + c.done.Do(func() { + c.finalized.Store(true) + _ = c.file.Close() + _ = os.Remove(c.fileName) + c.cloner.releaseCapture() + }) +} + +func (c *trafficCloneBodyCapture) fail(err error) { + c.done.Do(func() { + c.finalized.Store(true) + _ = c.file.Close() + _ = os.Remove(c.fileName) + for _, target := range c.targets { + labels := trafficCloneLabels(c.source, target) + trafficCloneRequests.With(trafficCloneResultLabels(labels, "error")).Inc() + log.Errorf("%s: cannot write traffic clone capture for target %q: %s", c.source.userName, target.name, err) + } + c.cloner.releaseCapture() + }) +} + +func (c *trafficCloneBodyCapture) skip(reason string) { + c.done.Do(func() { + c.finalized.Store(true) + _ = c.file.Close() + _ = os.Remove(c.fileName) + for _, target := range c.targets { + labels := trafficCloneLabels(c.source, target) + trafficCloneRequests.With(trafficCloneResultLabels(labels, "skipped")).Inc() + log.Debugf("%s: traffic clone to target %q skipped: %s", c.source.userName, target.name, reason) + } + c.cloner.releaseCapture() + }) +}