Add snapshot WebSocket API #583

seanlinsley · 2024-08-15T16:12:51Z

This introduces a persistent WebSocket connection for all server communication. Additional details:

If unable to start a WebSocket connection, the collector falls back to the legacy API.
grant_logs is still used for every log snapshot, since it's needed to generate a S3 upload URL. This will be deprecated in a separate PR.
All messages are zlib-compressed protobuf. The collector sends snapshots to the server, and the server sends ServerMessages.
The existing state.GrantConfig struct has been replaced with the same structure inside the protobuf, so the server and collector stay in sync.
For future use, ServerMessage includes a Pause message that will let us disable duplicate collector instances for the same server, and an ExplainRun message for the Explain Upload feature that's currently in-progress.

runner/websocket.go

state/state.go

keiko713

Still haven't tested by myself yet but looks pretty straight forward and promising.
Left some comments for things I noticed so far. I'm planning to do more review next week 👀

logs/querysample/tracing.go

runner/full.go

runner/websocket.go

Since this is a new connection the server doesn't know the collector was previously paused, so instead wait for the server to re-pause this collector.

runner/full.go

runner/logs.go

runner/websocket.go

state/state.go

msakrejda

Works as advertised, and reconnects automatically, but I did run into a panic on Ctrl+C:

2024/08/16 14:03:51 V [server1] Uploading snapshot to websocket
^C2024/08/16 14:03:58 I Exiting...
2024/08/16 14:03:58 V Log file fsnotify watcher received stop signal
2024/08/16 14:03:58 E [server1] Error closing websocket: close tcp 127.0.0.1:38368->127.0.0.1:5200: use of closed network connection
2024/08/16 14:03:58 V Stopping log tail for /var/log/postgresql/postgresql-15-main.log (stop requested)
2024/08/16 14:03:58 V Log file fsnotify watcher received stop signal
2024/08/16 14:03:58 V Stopping log tail for /var/log/postgresql/postgresql-16-main.log (stop requested)
panic: runtime error: invalid memory address or nil pointer dereference
[signal SIGSEGV: segmentation violation code=0x1 addr=0x0 pc=0xe487ca]

goroutine 9032 [running]:
github.com/gorilla/websocket.(*Conn).Close(...)
	/home/maciek/code/collector/vendor/github.com/gorilla/websocket/conn.go:352
github.com/pganalyze/collector/runner.connect.func1()
	/home/maciek/code/collector/runner/websocket.go:63 +0x16a
created by github.com/pganalyze/collector/runner.connect in goroutine 126
	/home/maciek/code/collector/runner/websocket.go:59 +0x78f

main.go

runner/websocket.go

runner/full.go

In passing, increase the default websocket buffer size from 4KB to 10KB

…rant config

seanlinsley · 2024-09-12T22:49:10Z

It looks like the test suite failed because it's trying to open a WebSocket connection to the production servers, but that side of the code hasn't been merged yet.

… full snapshots)

msakrejda

I tested this with https://github.com/pganalyze/pganalyze/pull/4234 and it worked fine. I left some minor comments on the code, but it looks fine, too.

runner/websocket.go

lfittl

Overall looks good - a few minor things pointed out, but seems close to be ready to merge.

state/state.go

runner/websocket.go

output/compact.go

protobuf/server_message.proto

runner/activity.go

runner/full.go

gorilla/websocket had to be updated: gorilla/websocket#877

seanlinsley · 2024-09-25T16:46:25Z

Okay, I believe all of the review comments have been addressed. I also fixed the issue where quitting the collector through Control-C showed a websocket warning.

This ensures we don't try connecting to the websocket while doing a log test

seanlinsley · 2024-09-30T22:57:30Z

main.go

@@ -397,7 +400,7 @@ func main() {
 		CollectSystemInformation: !noSystemInformation,
 		StateFilename:            stateFilename,
 		WriteStateUpdate:         (!dryRun && !dryRunLogs && !testRun) || forceStateUpdate,
-		ForceEmptyGrant:          dryRun || dryRunLogs || benchmark,
+		ForceEmptyGrant:          dryRun || dryRunLogs || testRunLogs || benchmark,


This was needed for the test suite to pass (otherwise the test suite would try and fail to establish a websocket connection). It's not clear from the main.go documentation what the difference is between --test-logs and --dry-run-logs. Maybe one of them should be dropped?

Hmm, so I think the different is:

--test-logs is supposed to get a log grant, and then either do a log download or log tail

--dry-run-logs is supposed to work without an API server, so skips the grant

However, I'm not sure if the distinction matters.

Really the main use case is --test-logs, which we document, to verify whether the log setup is correct. I suppose doing a log grant could be beneficial for Enterprise environments (since I think it would trigger the server side to check for correct LOG_ENCRYPTION_KEY / bucket setup?). Though that's soon to be removed anyway.

I'm not sure what the right course of action is here. Is there a blocker?

I don't think this is a blocker. We might want to remove --dry-run-logs after #597 is merged, though.

FWIW, I also don't think this is a blocker. I haven't seen much benefit of --dry-run-logs, so I'm good with dropping it.

That said, the regular --dry-run is useful though (it helps get clarity on what data the collector sent), and I've used that a few times for debugging with a customer, so we shouldn't be messing with that, I think :)

Sorry, I didn't explain my reasoning here. I'm thinking that there's little benefit to actually submitting test snapshots to the API, so in dropping --dry-run-logs my intention was to make --test-logs a dry run command. Then we only need to contact the API for a top-level --test in order to check for a valid grant config.

lfittl · 2024-10-01T13:47:57Z

* If unable to start a WebSocket connection, the collector falls back to the legacy API.

Not sure where exactly, but it appears there is a problem with the fallback, at least in my local test right now (on the latest tip of this branch) I got the following error, running against the regular server without websocket support:

% ./pganalyze-collector --test
2024/10/01 09:42:57 I Running collector test with pganalyze-collector 0.58.0
2024/10/01 09:42:57 W [server2] Checking log_line_prefix: database (%d) and user (%u) not found: pganalyze will not be able to correctly classify some log lines
2024/10/01 09:42:57 I [server2] Testing statistics collection...
2024/10/01 09:42:57 I [server1] Testing statistics collection...
2024/10/01 09:42:57 W [server1] Error starting websocket: websocket: bad handshake &{404 Not Found 404 HTTP/1.1 1 1 map[Content-Length:[36] Content-Type:[text/plain; charset=utf-8] Date:[Tue, 01 Oct 2024 13:42:57 GMT]] {0x140005d40c0} 36 [] false false map[] 0x14000764d00 <nil>}
2024/10/01 09:42:57 W [server2] Error starting websocket: websocket: bad handshake &{404 Not Found 404 HTTP/1.1 1 1 map[Content-Length:[36] Content-Type:[text/plain; charset=utf-8] Date:[Tue, 01 Oct 2024 13:42:57 GMT]] {0x140004d2b70} 36 [] false false map[] 0x14000578400 <nil>}
2024/10/01 09:42:57 I [server2]   Test collection successful for PostgreSQL 16.3 on aarch64-darwin, compiled by clang-14.0.0, 64-bit
2024/10/01 09:42:57 E [server2] Error uploading snapshot: Error - can't upload without valid S3 grant
2024/10/01 09:42:57 E [server2] Could not process server: Error - can't upload without valid S3 grant
2024/10/01 09:42:58 I [server1]   Test collection successful for PostgreSQL 16.3 (Debian 16.3-1.pgdg120+1) on aarch64-unknown-linux-gnu, compiled by gcc (Debian 12.2.0-14) 12.2.0, 64-bit
2024/10/01 09:42:58 E [server1] Error uploading snapshot: Error - can't upload without valid S3 grant
2024/10/01 09:42:58 E [server1] Could not process server: Error - can't upload without valid S3 grant
2024/10/01 09:42:58 I [server2] Testing activity snapshots...
2024/10/01 09:42:58 I [server1] Testing activity snapshots...
2024/10/01 09:42:58 I [server2]   Test snapshot successful
2024/10/01 09:42:58 I [server1]   Test snapshot successful

(note the "Error - can't upload without valid S3 grant", which I don't get when running on main)

seanlinsley · 2024-10-01T15:24:34Z

That was actually a bug from removing an existing write to server.Grant that you asked for in the previous review, which I've added back in 6f2bb88. There's no harm in potentially writing to server.Grant multiple times, so I think it's fine to leave the code as-is.

lfittl · 2024-10-01T16:18:48Z

That was actually a bug from removing an existing write to server.Grant that you asked for in the previous review, which I've added back in 6f2bb88. There's no harm in potentially writing to server.Grant multiple times, so I think it's fine to leave the code as-is.

Ah, makes sense - now we know why its needed :)

msakrejda

Assuming the ForceEmptyGrant change is okay (comment inline), this looks good.

msakrejda · 2024-10-01T19:26:39Z

main.go

@@ -397,7 +400,7 @@ func main() {
 		CollectSystemInformation: !noSystemInformation,
 		StateFilename:            stateFilename,
 		WriteStateUpdate:         (!dryRun && !dryRunLogs && !testRun) || forceStateUpdate,
-		ForceEmptyGrant:          dryRun || dryRunLogs || benchmark,
+		ForceEmptyGrant:          dryRun || dryRunLogs || testRunLogs || benchmark,


I'm not sure what the right course of action is here. Is there a blocker?

Add snapshot WebSocket API

6619535

seanlinsley commented Aug 15, 2024

View reviewed changes

runner/websocket.go Outdated Show resolved Hide resolved

seanlinsley commented Aug 15, 2024

View reviewed changes

state/state.go Outdated Show resolved Hide resolved

seanlinsley commented Aug 15, 2024

View reviewed changes

state/state.go Outdated Show resolved Hide resolved

seanlinsley requested a review from a team August 15, 2024 16:25

keiko713 reviewed Aug 16, 2024

View reviewed changes

logs/querysample/tracing.go Show resolved Hide resolved

runner/full.go Show resolved Hide resolved

runner/websocket.go Outdated Show resolved Hide resolved

runner/websocket.go Show resolved Hide resolved

seanlinsley added 3 commits August 16, 2024 10:39

Fix nil pointer errors

98b6cdd

Print snapshot submission logs similar to the legacy API

8e0234d

When connecting to a new websocket, ensure collector is no longer paused

e8d12d4

Since this is a new connection the server doesn't know the collector was previously paused, so instead wait for the server to re-pause this collector.

msakrejda reviewed Aug 16, 2024

View reviewed changes

review updates

679321d

keiko713 reviewed Aug 20, 2024

View reviewed changes

main.go Show resolved Hide resolved

runner/websocket.go Outdated Show resolved Hide resolved

runner/full.go Show resolved Hide resolved

seanlinsley added 10 commits August 20, 2024 09:48

update error handling

4202ed5

call the websocket's cancel function so it isn't leaked

8d076e6

update log level

6d43ba2

update websocket URL, protocol header, retry rate

30db024

Fix data races using atomic.Pointer

195a0d2

Merge branch 'main' into websocket

297b1e1

Fix test errors

ab8a23b

Add proxy support

2e019f4

In passing, increase the default websocket buffer size from 4KB to 10KB

Fix nil pointer error

825551d

Wrap Grant in atomic.Pointer, so JSON deserialization works for the g…

f48f263

…rant config

seanlinsley requested review from keiko713 and msakrejda September 12, 2024 22:19

Add test run output

54df371

seanlinsley added 3 commits September 23, 2024 12:08

Print summary of compact snapshots submitted to websocket

5ec4600

Don't print warning when server closes websocket connection normally

c676733

Change pause log to verbose for activity snapshots (still warning for…

6f8276b

… full snapshots)

msakrejda reviewed Sep 23, 2024

View reviewed changes

runner/websocket.go Outdated Show resolved Hide resolved

runner/websocket.go Outdated Show resolved Hide resolved

runner/websocket.go Show resolved Hide resolved

runner/websocket.go Show resolved Hide resolved

seanlinsley added 3 commits September 23, 2024 18:46

Close websocket connection if message decompression fails

218f824

Merge branch 'main' into websocket

793f91a

Fix memory leak caused by using defer in a loop

fbbeb42

lfittl reviewed Sep 24, 2024

View reviewed changes

seanlinsley added 7 commits September 25, 2024 10:24

small changes

4f7da1b

Use server_url from grant instead of test endpoint

81eee06

Don't assume websocket has provided valid grant yet

4558065

Only set grant during full snapshot when websocket is not active

b8fac4e

Silence websocket warning when quitting the collector process

f3d73da

gorilla/websocket had to be updated: gorilla/websocket#877

Tell server if this websocket is part of a test run

6e9f6ea

Add URL parse error handling

529acc2

seanlinsley added 3 commits September 26, 2024 12:53

Remove pause reason

06f316a

Don't try connecting to websocket during dry runs

2f9565a

Force empty grant for --test-logs

6f4ebf0

This ensures we don't try connecting to the websocket while doing a log test

seanlinsley commented Sep 30, 2024

View reviewed changes

Store grant on server in fallback mode to legacy API

6f2bb88

msakrejda approved these changes Oct 1, 2024

View reviewed changes

seanlinsley mentioned this pull request Oct 1, 2024

Include log text in snapshots #597

Merged

2 tasks

seanlinsley merged commit d7a5290 into main Oct 3, 2024
3 checks passed

seanlinsley deleted the websocket branch October 3, 2024 00:36

seanlinsley mentioned this pull request Oct 7, 2024

Improve test run error reporting for WebSocket #608

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add snapshot WebSocket API #583

Add snapshot WebSocket API #583

seanlinsley commented Aug 15, 2024 •

edited

Loading

keiko713 left a comment

msakrejda left a comment

seanlinsley commented Sep 12, 2024

msakrejda left a comment

lfittl left a comment

seanlinsley commented Sep 25, 2024

seanlinsley Sep 30, 2024 •

edited

Loading

lfittl Oct 1, 2024

msakrejda Oct 1, 2024

seanlinsley Oct 1, 2024

lfittl Oct 2, 2024 •

edited

Loading

seanlinsley Oct 5, 2024 •

edited

Loading

lfittl commented Oct 1, 2024

seanlinsley commented Oct 1, 2024

lfittl commented Oct 1, 2024

msakrejda left a comment

msakrejda Oct 1, 2024

Add snapshot WebSocket API #583

Add snapshot WebSocket API #583

Conversation

seanlinsley commented Aug 15, 2024 • edited Loading

keiko713 left a comment

Choose a reason for hiding this comment

msakrejda left a comment

Choose a reason for hiding this comment

seanlinsley commented Sep 12, 2024

msakrejda left a comment

Choose a reason for hiding this comment

lfittl left a comment

Choose a reason for hiding this comment

seanlinsley commented Sep 25, 2024

seanlinsley Sep 30, 2024 • edited Loading

Choose a reason for hiding this comment

lfittl Oct 1, 2024

Choose a reason for hiding this comment

msakrejda Oct 1, 2024

Choose a reason for hiding this comment

seanlinsley Oct 1, 2024

Choose a reason for hiding this comment

lfittl Oct 2, 2024 • edited Loading

Choose a reason for hiding this comment

seanlinsley Oct 5, 2024 • edited Loading

Choose a reason for hiding this comment

lfittl commented Oct 1, 2024

seanlinsley commented Oct 1, 2024

lfittl commented Oct 1, 2024

msakrejda left a comment

Choose a reason for hiding this comment

msakrejda Oct 1, 2024

Choose a reason for hiding this comment

seanlinsley commented Aug 15, 2024 •

edited

Loading

seanlinsley Sep 30, 2024 •

edited

Loading

lfittl Oct 2, 2024 •

edited

Loading

seanlinsley Oct 5, 2024 •

edited

Loading