Reduce message cloning operations #190

zainkabani · 2022-10-14T15:11:01Z

This PR leverages client and server buffers to read messages onto and flush messages from.

This reduces the number of allocations performed by pgcat as seen by these stats after pgbench run on the main branch compared to this branch.

pgbench with new client per query results in a ~37% reduction in allocations
pgbench -t 1000 -c 16 -j 2 --protocol extended -C

Before:

Stats: Stats {
    allocations: 5617437,
    deallocations: 5614634,
    reallocations: 130699,
    bytes_allocated: 673340233,
    bytes_deallocated: 672221777,
    bytes_reallocated: 3760681,
}

After:

Stats: Stats {
    allocations: 3505408,
    deallocations: 3502605,
    reallocations: 130702,
    bytes_allocated: 628430295,
    bytes_deallocated: 627311847,
    bytes_reallocated: 3760728,
}

pgbench with same client for pgbench session results in a ~53% reduction in allocations:
pgbench -t 1000 -c 16 -j 2 --protocol extended

Before:

Stats: Stats {
    allocations: 3874108,
    deallocations: 3871232,
    reallocations: 2880,
    bytes_allocated: 158461974,
    bytes_deallocated: 157277508,
    bytes_reallocated: 797417,
}

After:

Stats: Stats {
    allocations: 1792054,
    deallocations: 1789265,
    reallocations: 2804,
    bytes_allocated: 113265426,
    bytes_deallocated: 112172240,
    bytes_reallocated: 780508,
}

levkk · 2022-10-14T17:34:33Z

src/admin.rs

    res.put_u8(b'I');

-    write_all_half(stream, res).await
+    write_all_half(stream, &res).await


I don't think this changes anything. res is deallocated at the end of this block anyway, so borrowing it doesn't do anything, you might as well pass ownership to write_all_half(). Passing ownership != copying.

if the function doesn't need the explicit object then it's better to have it accept the reference so it can stay more flexible

Another reason we're using reference here is because we don't want to clone or pass ownership of the buffer to the send operation so we pass a reference instead.

levkk · 2022-10-14T19:26:32Z

src/server.rs

    /// in order to receive all data the server has to offer.
    pub async fn recv(&mut self) -> Result<BytesMut, Error> {
+        // Our server response buffer. We buffer data before we give it to the client.
+        let mut message_buffer = BytesMut::with_capacity(8196);


This will allocate a new buffer every time we send a query, this will be much slower than it is now.

which is more expensive, cloning or allocating?
cloning does both so if anything is is equivalent

Why not just pass self.buffer[..] here instead avoiding both problems? I just know that in network programming, the buffer is never reallocated and is re-used instead, so we should make sure to do that as well.

To answer your question, I am not sure. My gut feeling makes me thing allocating because it'll be allocated on the heap, but maybe cloning does the same thing in this case.

the problem here is we're writing into the server's buffer not the client's send buffer.
also self.buffer[..] returns a different type. The problem is that the buffer isn't shared between the client and server so we can't really populate and flush. Also rust's ownership model makes it difficult too. I can imagine a re-architecture that solves this better but for now this change maintains parity if anything

levkk · 2022-10-17T18:46:34Z

src/client.rs

        // Internal buffer, where we place messages until we have to flush
        // them to the backend.
-        let mut message_buffer = BytesMut::with_capacity(8196);
+        let mut client_message_buffer = BytesMut::with_capacity(8196);


Both client and server have their own buffers already, both self.buffer. I unfortunately don't understand what you're trying to accomplish here.

This idea came from your earlier comment about how in networking we have a buffer that we usually don't allocate and deallocate frequently. When writing to the client the server needs to populate the buffer that the client uses to send messages, which is what the server_message_buffer is. We were previously using the server's own buffer and then cloning it, which is an allocation and data copy.

Yeah...this ends up being the same thing I think. Either we write to server.buffer or this new buffer, we have to move it out of the server buffer and clear it, and since we're clearing it, we need to clone whatever was in it in the first place.

I think to achieve what you want, we need to use slices in a very clever way. Basically when a new message comes in, we make read_message read that data directly into the clients buffer at the offset where we know there is no data and to ensure there is no empty bytes in between messages. We then pass the clients buffer around as a reference and whoever needs to parse the query / message inside it can do so without modifying the buffer. Finally we pass it to the server send function (borrow again) and write its contents directly to the server socket. Finally, we clear the buffer. No buffer copies, just one single copy when reading the clients message and one copy when writing it to the Postgres backend.

This approaches C-style programming honestly, and it's not straight forward. I think maybe the bytes crate has some handy methods for us to use? Worth a look, maybe.

dat2 · 2022-11-16T21:01:59Z

src/client.rs

    /// Internal buffer, where we place messages until we have to flush
    /// them to the backend.
-    buffer: BytesMut,
+    client_message_buffer: BytesMut,


we should keep the name as buffer as its inside a Client struct

dat2 · 2022-11-16T21:03:17Z

src/client.rs

                // to when we get the S message
                'P' | 'B' | 'D' | 'E' => {
-                    self.buffer.put(&message[..]);
+                    // client_message_buffer.put(&message[..]);


can we remove the commented code?

dat2 · 2022-11-16T21:08:14Z

src/messages.rs

-    bytes.put_slice(&buf);
+    buffer.put_u8(code);
+    buffer.put_i32(len);
+    buffer.put_slice(&buf);


Suggested change

buffer.put_slice(&buf);

buffer.put_slice(&buffer);

i think this would allow us to get rid of the let mut buf allocation, cc @levkk correct me.

Removing the creation of this buffer to read the message body was a huge lift! Great idea :D

dat2 · 2022-11-16T21:11:10Z

src/client.rs

                    // Admin clients ignore shutdown.
                    else {
-                        read_message(&mut self.read).await?
+                        read_message(&mut self.read, &mut self.client_message_buffer).await?


I think we should rename read_message to read_message_into_buffer now. the old implementation returned a buffer, and now it writes values into the buffer. i don't have a strong opinion on this, but @levkk what do you think?

It would be nice to keep both for different use cases, I agree.

dat2 · 2022-11-16T21:20:31Z

src/query_router.rs

-        let len = buf.get_i32() as usize;
-        let query = String::from_utf8_lossy(&buf[..len - 5]).to_string(); // Ignore the terminating NULL.
+        let _len = message_cursor.get_i32() as usize;
+        let query = message_cursor.read_string().unwrap();


can we keep query as a &str? do we need to convert it to an owned String

I don't recall, but in all likelihood, we can keep it as &str.

Can't easily return &str, open to suggestions on how to do that though

levkk · 2022-11-18T22:17:20Z

src/lib.rs

+}
+
+impl BytesMutReader for Cursor<&BytesMut> {
+    fn read_string(&mut self) -> Result<String, Error> {


What if the query contains a null byte? I think we need to be careful here and read the exact amount Postrgres is telling us is in the message.

Yeah that's a fair call out, however this function should only be used to read strings explictly from things like the query packet, or parameters and the startupparameters packet. Postgres specifically doesn't allow null bytes because they're used as special characters in the protocol.

some info about that here, https://stackoverflow.com/questions/28813409/are-null-bytes-allowed-in-unicode-strings-in-postgresql-via-python#:~:text=Actually%2C%20this%20is%20only%20mentioned,types%20cannot%20store%20such%20bytes.

levkk · 2022-11-18T22:18:01Z

src/messages.rs

+
+    match stream
+        .read_exact(
+            &mut buffer[starting_point + mem::size_of::<u8>() + mem::size_of::<i32>()


You're sometimes hardcoding the length of the integer (4 bytes) and sometimes using these functions. Might be best to stick to one or the other.

levkk · 2022-11-18T22:20:07Z

src/query_router.rs

            // Query
            'Q' => {
-                let query = String::from_utf8_lossy(&buf[..len - 5]).to_string();
+                let query = message_cursor.read_string().unwrap();


Here we should really read the exact amount we know is in the message instead of reading until a NULL byte.

…in read_message

… and clear server buffer Fixes incorrect log variable bug

zainkabani added 2 commits October 14, 2022 11:10

initial commit

d7bf2de

fix typo

b42c33c

levkk reviewed Oct 14, 2022

View reviewed changes

zainkabani added 2 commits October 14, 2022 14:28

use cursor for parse params instead of bytesmut

3ae2953

undo

33a0cad

zainkabani marked this pull request as ready for review October 14, 2022 18:43

levkk reviewed Oct 14, 2022

View reviewed changes

zainkabani added 2 commits October 17, 2022 10:52

Update to use a dedicated server message buffer

65bd10d

fmt

09db31d

levkk reviewed Oct 17, 2022

View reviewed changes

zainkabani added 7 commits November 14, 2022 20:51

Read message directly onto buffer instead of new bytesmut

8b47f32

Merge branch 'main' into zain/reduce-cloning-operations

d6c8271

remove commented code

1589548

Move server buffer to server object

e2b2cb0

Move client and server message buffers to attribute of object

8d68d22

Remove clear buffer function since public already

f019e28

Remove commented code

4cb38b8

dat2 reviewed Nov 16, 2022

View reviewed changes

Remove cloning operation for query router

8f3fc25

dat2 reviewed Nov 16, 2022

View reviewed changes

zainkabani added 7 commits November 16, 2022 16:26

Rename server and client buffers to buffer

728f6cf

Fix bug

ee19dbc

Rename read_messages to read_messages_into_buffer

027011d

fmt

f0ceb8b

Rename to message_buffer to be more explicit

021a108

Read message body directly onto buffer

24d36af

Fix bug

8a58f8b

zainkabani added 6 commits November 16, 2022 18:28

fmt

089348b

Unused try_into

ecb9627

use read_exact instead of read to guarantee number of bytes read

a3e3801

Remove length check

bece3b9

fmt

54ed58e

Merge branch 'main' into zain/reduce-cloning-operations

ada475e

levkk reviewed Nov 18, 2022

View reviewed changes

zainkabani added 14 commits November 21, 2022 11:34

Add comment to read_string function and uses sizeof for byte lengths …

ccde8f0

…in read_message

fmt

6c5b5df

Refactor reading query in infer role

7bf9503

Creates send server message to client function to both send to client…

5f8473c

… and clear server buffer Fixes incorrect log variable bug

Merge branch 'main' into zain/reduce-cloning-operations

0fa0214

Merge branch 'main' into zain/reduce-cloning-operations

2986f06

Ensures buffer is cleared before it checked in

382e0b8

revert?

459e5d4

try again

ae4078e

Merge branch 'main' into zain/reduce-cloning-operations

08f179c

Merge branch 'main' into zain/reduce-cloning-operations

6ad7564

Update comments

2a84549

fmt

fb891fc

fix

9937aba

Reduce message cloning operations #190

Are you sure you want to change the base?

Reduce message cloning operations #190

Uh oh!

Conversation

zainkabani commented Oct 14, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zainkabani Oct 14, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zainkabani commented Oct 14, 2022 •

edited

Loading

zainkabani Oct 14, 2022 •

edited

Loading