feat(sink): support async for bigquery sink #17488

xxhZs · 2024-06-27T10:20:01Z

I hereby agree to the terms of the RisingWave Labs, Inc. Contributor License Agreement.

What's changed and what's your intention?

in #17383 (comment)
We found that waiting to write data can block for a long time, affecting throughput.
So we support async for bigquery sink

bench with this pr:
avg: 108025 rows/s
p90: 116736 rows/s
p95: 116736 rows/s
p99: 118784 rows/s

Checklist

I have written necessary rustdoc comments
I have added necessary unit tests and integration tests
I have added test labels as necessary. See details.
I have added fuzzing tests or opened an issue to track them. (Optional, recommended for new SQL features Sqlsmith: Sql feature generation #7934).
My PR contains breaking changes. (If it deprecates some features, please create a tracking issue to remove them in the future).
All checks passed in ./risedev check (or alias, ./risedev c)
My PR changes performance-critical code. (Please run macro/micro-benchmarks and show the results.)

My PR contains critical fixes that are necessary to be merged into the latest release. (Please check out the details)

Documentation

My PR needs documentation updates. (Please use the Release note section below to summarize the impact on users)

Release note

If this PR includes changes that directly affect users or other significant modifications relevant to the community, kindly draft a release note to provide a concise summary of these changes. Please prioritize highlighting the impact these changes will have on users.

recopver

wenym1

Actually we can avoid spawning a tokio task to poll the response stream. We can implement our own BigQueryLogSinker, which polls the log_reader and the response_stream with a select.

wenym1 · 2024-07-04T07:02:49Z

src/connector/src/sink/big_query.rs

+        let (expect_offset, mut rx) = self.client.get_subscribe();
+        let future = Box::pin(async move {
+            loop {
+                match rx.recv().await {


I think instead you can create a oneshot channel for each AppendRowsRequest so that you won't need this shared broadcast channel. When the spawned worker knows that a request is handled, it notifies the rx with the oneshot tx, and here the future can be simply the rx.

wenym1 · 2024-07-04T07:04:18Z

src/connector/src/sink/big_query.rs

+        let append_req = AppendRowsRequest {
+            write_stream: write_stream.clone(),
+            offset: None,
+            trace_id: Uuid::new_v4().hyphenated().to_string(),


Can you tracked the trace_id locally instead of incrementing the self.offset? I think in the response stream, it should return the same trace_id in the corresponding response.

Tried it, but there is no ID in the returned message

fix

src/connector/src/sink/big_query.rs

svae add fix

xxhZs · 2024-07-22T05:45:19Z

Since each row of bg can't exceed 10MB, new logic is added to split the size of each row

stdrc · 2024-07-26T07:29:10Z

src/connector/src/sink/big_query.rs

+        while resp_num > 1 {
+            self.offset_queue.push_back(None);
+            resp_num -= 1;
+        }
+        self.offset_queue.push_back(Some(offset));


Why inserting resp_num - 1 Nones and 1 Some(offset)? Could you please add some comments here?

Because if a chunk is too large, we will split it into multiple async tasks, so we need to wait until all the tasks are finished before we can truncate, here none means that he is only the middle async task is completed, to the chunk before you can truncate.
And i will add comments.

stdrc · 2024-07-26T07:30:03Z

src/connector/src/sink/big_query.rs


-fn default_retry_times() -> usize {
-    5
+    pub async fn wait_next_offset(&mut self) -> Result<Option<TruncateOffset>> {


I guess async fn next_offset is enough to indicate that the function call may need to be "waited".

stdrc · 2024-07-26T07:32:48Z

src/connector/src/sink/big_query.rs

+        if let Some(Some(TruncateOffset::Barrier { .. })) = self.offset_queue.front() {
+            return Ok(self.offset_queue.pop_front().unwrap());
+        }
+        self.resp_stream
+            .next()
+            .await
+            .ok_or_else(|| SinkError::BigQuery(anyhow::anyhow!("end of stream")))??;
+        self.offset_queue.pop_front().ok_or_else(|| {
+            SinkError::BigQuery(anyhow::anyhow!(
+                "should have pending chunk offset when we receive new response"
+            ))
+        })


Why we can only directly pop and return if the front of queue is TruncateOffset::Barrier? Is it possible to have other variants of TruncateOffset in the queue? And what happen if there are?

Could you please add some comments here?

Since we don't go to the bg sink to create an async task when we receive a barrier, there's no need to wait here for the
And i will add comments.

Since we don't go to the bg sink to create an async task when we receive a barrier, there's no need to wait here for the

Can't quite get it. Could you elaborate in proper English?

There are 3 cases:

Some(barrier), since we won't be sending a barrier to bigquery, there is no need to wait for the. But we will record the barriers in the log store, so we need to truncate.

Some(chunk), since we send a chunk to bigquery and record it in the log store, so we need to truncate and wait for resp.

None, It means that we split a large chunk of RW into multiple requests for bigquery (for this chunk, in offset_queue, we set it to none , none ... chunk). So we need to wait for all the none and chunk resps and truncate in Some(chunk)

stdrc · 2024-07-26T07:39:58Z

src/connector/src/sink/big_query.rs

+}
+pub struct BigQueryLogSinker {
+    writer: BigQuerySinkWriter,
+    bigquery_future_manager: BigQueryFutureManager,


Why not future_manager? we don't need to add so many prefixes in private fields/types. And I guess it may not be so suitable to call it "manager" here.

Also, do we really need a BigQueryFutureManager type here? It seems don't "manage" anything, we still directly access bigquery_future_manager.offset_queue below😂

For async sink, we have a FutureManager to manage the asyn tasks and truncate our log store chunk after the async tasks are done. this BigQueryFutureManager has the same tasks as FutureManager, just customized for bg sink

wenym1 · 2024-08-26T09:37:48Z

src/connector/src/sink/big_query.rs

-        })
+        let mut client = conn.conn();
+
+        let (tx, rx) = mpsc::channel(BIGQUERY_SEND_FUTURE_BUFFER_MAX_SIZE);


We can use bounded channel here since we already limit the queue size in the select!

wenym1 · 2024-08-26T09:42:31Z

src/connector/src/sink/big_query.rs

+        .map_err(|e| SinkError::BigQuery(e.into()))?
+        .into_inner();
+    loop {
+        if let Some(append_rows_response) = resp_stream


When we reach the end of the resp_stream, we will keep yielding (). This is not correct. Instead of using if let Some(...), we'd better just resp_stream.message().await.ok_or_else(|| ...) to turn None into an end-of-stream error.

It doesn't return () at the end of the stream, after receiving every reply, if there are no errors, it returns a meaningless (),
The goal is to return that a message was received with no errors

I know the purpose of returning a meaningless (). What I meant was that, when resp_stream.message() returns None, we should break the loop with an EndOfStream error, or just simply break the loop, instead of ignoring it.

In current code, after resp_stream returns None, an () will still be yielded, and later when polled again, and resp_stream returns None again, and an () is still yielded, and this will be repeated endlessly over and over again, while the external code is not aware that the response stream has actually stopped.

wenym1 · 2024-08-26T09:58:26Z

src/connector/src/sink/big_query.rs

+        if let Some(Some(TruncateOffset::Barrier { .. })) = self.offset_queue.front() {
+            return Ok(self.offset_queue.pop_front().unwrap());
+        }
+        self.resp_stream


Here the correctness depends on that the number of non-barrier items (either None or Some(TruncateOffset::Chunk)) in the queue is the same as the inflight request. However, in the implementation, it's possible to have 0 resp_num when we write a chunk. For this chunk, there is no inflight request, but it will have an item in the queue, which causes inconsistency between the number of queue items and inflight requests.

I think instead of using none to represent that a chunk is split into multiple requests, we'd better store (TruncateOffset, remaining_resp_num) as the queue item. Code will be like the following

if let Some((offset, remaining_resp_num)) = self.offset_queue.front_mut() { if *remaining_resp_num == 0 { return Ok(self.offset_queue.pop_front().unwrap().0); } while *remaining_resp_num > 0 { self.resp_stream.next().await...??; *remaining_resp_num -= 1; } } else { return pending().await; }

save save

wenym1

Rest LGTM

wenym1 · 2024-09-02T04:36:40Z

src/bench/sink_bench/sink_option.yml

@@ -100,8 +100,8 @@ Starrocks:
 BigQuery:


Revert the change.

wenym1 · 2024-09-02T04:36:59Z

src/connector/src/sink/big_query.rs

-fn default_max_batch_rows() -> usize {
-    1024
+struct BigQueryFutureManager {
+    // `offset_queue` holds the Some corresponding to each future.


Please update the comment.

xxhZs added 5 commits June 27, 2024 15:46

support

4baebce

support

df1a80b

remove option

eed3cac

Merge branch 'main' into xxh/bg-async

7df4521

fmt

7a311dc

xxhZs requested a review from a team as a code owner June 27, 2024 10:36

github-actions bot added the type/feature label Jun 27, 2024

recovery cargo lock

e46e816

recopver

xxhZs force-pushed the xxh/bg-async branch from 921dc3d to e46e816 Compare June 27, 2024 10:39

fix ci

f097e35

xxhZs requested a review from wenym1 July 1, 2024 08:17

wenym1 reviewed Jul 4, 2024

View reviewed changes

xxhZs added 5 commits July 4, 2024 18:41

refactor

9c94778

Merge branch 'main' into xxh/bg-async

37101b1

fix

ca9e599

save

09b645d

refa

02352ac

fix

xxhZs force-pushed the xxh/bg-async branch from 00f7738 to 02352ac Compare July 8, 2024 10:43

xxhZs added 2 commits July 8, 2024 19:02

Merge branch 'main' into xxh/bg-async

568773a

fix ci

f26bf06

xxhZs requested a review from wenym1 July 10, 2024 08:52

wenym1 reviewed Jul 16, 2024

View reviewed changes

src/connector/src/sink/big_query.rs Outdated Show resolved Hide resolved

xxhZs added 2 commits July 19, 2024 12:39

fix comm

181978c

fix ci

9cc98ca

svae add fix

fmt

05d3922

fuyufjh removed the request for review from a team July 23, 2024 08:03

stdrc reviewed Jul 26, 2024

View reviewed changes

xxhZs added 2 commits July 29, 2024 11:50

add doc

5721414

Merge branch 'main' into xxh/bg-async

aff5ac8

fix comm

ecd13d6

xxhZs mentioned this pull request Aug 14, 2024

Tracking: enable sink_decouple by default for all types of sink #17095

Closed

31 tasks

wenym1 reviewed Aug 26, 2024

View reviewed changes

save

485b0e7

save save

wenym1 approved these changes Sep 2, 2024

View reviewed changes

xxhZs and others added 2 commits September 2, 2024 14:49

fix comm

8f58c7f

Merge branch 'main' into xxh/bg-async

040d2da

xxhZs enabled auto-merge September 3, 2024 08:43

xxhZs added this pull request to the merge queue Sep 3, 2024

Merged via the queue into main with commit c0ce8a8 Sep 3, 2024
29 of 30 checks passed

xxhZs deleted the xxh/bg-async branch September 3, 2024 09:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(sink): support async for bigquery sink #17488

feat(sink): support async for bigquery sink #17488

xxhZs commented Jun 27, 2024

wenym1 left a comment

wenym1 Jul 4, 2024

wenym1 Jul 4, 2024

xxhZs Jul 8, 2024

xxhZs commented Jul 22, 2024

stdrc Jul 26, 2024

xxhZs Jul 26, 2024

stdrc Jul 26, 2024

stdrc Jul 26, 2024

xxhZs Jul 26, 2024

stdrc Jul 29, 2024 •

edited

Loading

xxhZs Jul 29, 2024 •

edited

Loading

stdrc Jul 26, 2024

xxhZs Jul 26, 2024

wenym1 Aug 26, 2024

wenym1 Aug 26, 2024

xxhZs Aug 28, 2024 •

edited

Loading

wenym1 Aug 29, 2024

wenym1 Aug 26, 2024

wenym1 left a comment

wenym1 Sep 2, 2024

wenym1 Sep 2, 2024

feat(sink): support async for bigquery sink #17488

feat(sink): support async for bigquery sink #17488

Conversation

xxhZs commented Jun 27, 2024

What's changed and what's your intention?

Checklist

Documentation

Release note

wenym1 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xxhZs commented Jul 22, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stdrc Jul 29, 2024 • edited Loading

Choose a reason for hiding this comment

xxhZs Jul 29, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xxhZs Aug 28, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wenym1 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stdrc Jul 29, 2024 •

edited

Loading

xxhZs Jul 29, 2024 •

edited

Loading

xxhZs Aug 28, 2024 •

edited

Loading