Skip to content

kgen: parallelize #9841

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jan 1, 2022
Merged

kgen: parallelize #9841

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
255da7b
ore/retry: create a separate RetryStream type
benesch Jan 1, 2022
6561cb0
ore/retry: rename Retry::retry to Retry::retry_async
benesch Jan 1, 2022
d6aab5e
ore/retry: add a synchronous Retry::retry operation
benesch Jan 1, 2022
d65340a
ore/retry: use newly-stabilized Duration::MAX constant
benesch Jan 1, 2022
f6a6c66
kgen: use exponential backoff to retry when queue is full
benesch Jan 1, 2022
1e444ab
kgen: parallelize
benesch Jan 1, 2022
fcb0f7d
benches/avro_ingest: use kgen directly rather than kafka-avro-generator
benesch Jan 1, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 28 additions & 9 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion demo/billing/src/mz.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ pub async fn validate_sink(
let count_input_view_query = format!("SELECT count(*) from {}", input_view);

Retry::default()
.retry(|_| async {
.retry_async(|_| async {
let count_check_sink: i64 = mz_client
.query_one(&*count_check_sink_query, &[])
.await?
Expand Down
139 changes: 120 additions & 19 deletions misc/python/materialize/benches/avro_ingest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
"""Ingest some Avro records, and report how long it takes"""

import argparse
import json
import os
import time
from typing import IO, NamedTuple
Expand Down Expand Up @@ -87,21 +88,14 @@ def main() -> None:
type=int,
help="Number of Avro records to generate",
)
parser.add_argument(
"-d",
"--distribution",
default="benchmark",
type=str,
help="Distribution to use in kafka-avro-generator",
)
args = parser.parse_args()

os.chdir(ROOT)
repo = mzbuild.Repository(ROOT)

wait_for_confluent(args.confluent_host)

images = ["kafka-avro-generator", "materialized"]
images = ["kgen", "materialized"]
deps = repo.resolve_dependencies([repo.images[name] for name in images])
deps.acquire()

Expand All @@ -114,18 +108,18 @@ def main() -> None:
)

docker_client.containers.run(
deps["kafka-avro-generator"].spec(),
deps["kgen"].spec(),
[
"-n",
str(args.records),
"-b",
f"{args.confluent_host}:9092",
"-r",
f"http://{args.confluent_host}:8081",
"-t",
"bench_data",
"-d",
args.distribution,
f"--num-records={args.records}",
f"--bootstrap-server={args.confluent_host}:9092",
f"--schema-registry-url=http://{args.confluent_host}:8081",
"--topic=bench_data",
"--keys=avro",
"--values=avro",
f"--avro-schema={VALUE_SCHEMA}",
f"--avro-distribution={VALUE_DISTRIBUTION}",
f"--avro-key-schema={KEY_SCHEMA}",
f"--avro-key-distribution={KEY_DISTRIBUTION}",
],
network_mode="host",
)
Expand Down Expand Up @@ -158,5 +152,112 @@ def main() -> None:
prev = print_stats(mz_container, prev, results_file)


KEY_SCHEMA = json.dumps(
{
"name": "testrecordkey",
"type": "record",
"namespace": "com.acme.avro",
"fields": [{"name": "Key1", "type": "long"}, {"name": "Key2", "type": "long"}],
}
)

KEY_DISTRIBUTION = json.dumps(
{
"com.acme.avro.testrecordkey::Key1": [0, 100],
"com.acme.avro.testrecordkey::Key2": [0, 250000],
}
)

VALUE_SCHEMA = json.dumps(
{
"name": "testrecord",
"type": "record",
"namespace": "com.acme.avro",
"fields": [
{"name": "Key1Unused", "type": "long"},
{"name": "Key2Unused", "type": "long"},
{
"name": "OuterRecord",
"type": {
"name": "OuterRecord",
"type": "record",
"fields": [
{
"name": "Record1",
"type": {
"name": "Record1",
"type": "record",
"fields": [
{
"name": "InnerRecord1",
"type": {
"name": "InnerRecord1",
"type": "record",
"fields": [
{"name": "Point", "type": "long"}
],
},
},
{
"name": "InnerRecord2",
"type": {
"name": "InnerRecord2",
"type": "record",
"fields": [
{"name": "Point", "type": "long"}
],
},
},
],
},
},
{
"name": "Record2",
"type": {
"name": "Record2",
"type": "record",
"fields": [
{
"name": "InnerRecord3",
"type": {
"name": "InnerRecord3",
"type": "record",
"fields": [
{"name": "Point", "type": "long"}
],
},
},
{
"name": "InnerRecord4",
"type": {
"name": "InnerRecord4",
"type": "record",
"fields": [
{"name": "Point", "type": "long"}
],
},
},
],
},
},
],
},
},
],
}
)

VALUE_DISTRIBUTION = json.dumps(
{
"com.acme.avro.testrecord::Key1Unused": [0, 100],
"com.acme.avro.testrecord::Key2Unused": [0, 250000],
"com.acme.avro.InnerRecord1::Point": [10000, 1000000000],
"com.acme.avro.InnerRecord2::Point": [10000, 1000000000],
"com.acme.avro.InnerRecord3::Point": [10000, 1000000000],
"com.acme.avro.InnerRecord4::Point": [10000, 10000000000],
}
)


if __name__ == "__main__":
main()
4 changes: 3 additions & 1 deletion src/coord/src/coord.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3913,7 +3913,9 @@ where
for (conn, slot_names) in replication_slots_to_drop {
// Try to drop the replication slots, but give up after a while.
let _ = Retry::default()
.retry(|_state| postgres_util::drop_replication_slots(&conn, &slot_names))
.retry_async(|_state| {
postgres_util::drop_replication_slots(&conn, &slot_names)
})
.await;
}
});
Expand Down
2 changes: 1 addition & 1 deletion src/dataflow/src/source/s3.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -312,7 +312,7 @@ async fn scan_bucket_task(
let mut continuation_token = None;
loop {
let response = Retry::default()
.retry(|_| {
.retry_async(|_| {
client
.list_objects_v2()
.bucket(&bucket)
Expand Down
2 changes: 1 addition & 1 deletion src/interchange/src/avro/schema.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -330,7 +330,7 @@ impl SchemaCache {
let ccsr_client = &self.ccsr_client;
let response = Retry::default()
.max_duration(Duration::from_secs(30))
.retry(|state| async move {
.retry_async(|state| async move {
let res = ccsr_client.get_schema_by_id(id).await;
match res {
Err(e) => {
Expand Down
2 changes: 2 additions & 0 deletions src/kafka-util/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,10 @@ anyhow = "1.0.52"
ccsr = { path = "../ccsr" }
chrono = { version = "0.4.0", default-features = false, features = ["std"] }
clap = "2.34.0"
crossbeam = "0.8.1"
futures = "0.3.19"
mz-avro = { path = "../avro" }
num_cpus = "1.13.1"
ore = { path = "../ore", features = ["network"] }
rand = "0.8.4"
rdkafka = { git = "https://github.com/fede1024/rust-rdkafka.git", features = ["cmake-build", "libz-static"] }
Expand Down
2 changes: 1 addition & 1 deletion src/kafka-util/src/admin.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ where
// created with the default number partitions, and not the number of
// partitions requested in `new_topic`.
Retry::default()
.retry(|_| async {
.retry_async(|_| async {
let metadata = client
.inner()
// N.B. It is extremely important not to ask specifically
Expand Down
Loading