ccl/sqlproxyccl: add postgres interceptors for message forwarding

[jaylim-crl]

Informs #76000.

This commit implements postgres interceptors, namely FrontendInterceptor and
BackendInterceptor, as described in the sqlproxy connection migration RFC.
These interceptors will be used as building blocks for the forwarder component
that we will be adding in a later PR. Since the forwarder component has not
been added, a simple proxy test (i.e. TestSimpleProxy) has been added to
illustrate how the frontend and backend interceptors can be used within the
proxy.

Release note: None

[cockroach-teamcity]

This change is 

[rafiss]

do you know if the FrontendInterceptor could also be used to intercept CancelRequests? (not needed now, but wondering for the future)

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @andy-kimball, @jaylim-crl, and @jeffswenson)

---

pkg/ccl/sqlproxyccl/interceptor/interceptor_test.go, line 23 at r1 (raw file):

	"github.com/stretchr/testify/require"
)

can this file also include a benchmark test?

[jeffswenson]

This is looking great. I like the overall structure and really appreciate the fact this is a standalone PR.

[jaylim-crl]

TFTRs!

do you know if the FrontendInterceptor could also be used to intercept CancelRequests? (not needed now, but wondering for the future)

As discussed over Slack, I originally thought that CancelRequest messages are just like any other messages, but it appears that it's part of the startup message, so no. These interceptors are meant to be used post-auth phase, where all messages must include 5 bytes for the headers (byte for type + int for length). In the future, if we'd like to combine both phases, we could modify the interceptor to support this.

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @andy-kimball, @jeffswenson, and @rafiss)

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 56 at r1 (raw file):

Previously, JeffSwenson (Jeff Swenson) wrote…

Does the interceptor return any recoverable errors? If all of the errors are final, I think we should remove .closed field. Tracking the .closed state adds some complexity to the code and I think it is reasonable to document the caller needs to throw away the interceptor if there is a read/write error.

The interceptor does not return any recoverable errors unless the caller wrapped src or dst to return custom errors. I've thought about this further, and I think I'll remove this. We'll just defer all of this logic to the caller. If the caller wraps any of those connections with custom errors, those errors should be handled accordingly. If this becomes an issue, we'll add this back.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 116 at r1 (raw file):

Previously, JeffSwenson (Jeff Swenson) wrote…

On 32 bit machines this cast can overflow. Most of the cases would be caught by the size<0 check below, but ~0+3, ~0+2, ~1, and ~0 would result in a weird state.

I'm not sure if it's worth handling this case given that we're already doing this on the DB side today:

cockroach/pkg/sql/pgwire/pgwirebase/encoding.go

Lines 154 to 156 in f1fcb83

	size := int(binary.BigEndian.Uint32(b.tmp[:]))
	// size includes itself.
	size -= 4

I'm thinking to leave this as-is unless you strongly feel about this.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 157 at r1 (raw file):

Previously, JeffSwenson (Jeff Swenson) wrote…

I recommend changing the API to return msg []byte instead of body []byte. That allows us to simplify the implementation of chunkReader.Next and removes a few instances of needing to add/subtract the header size.

Good point. I can do that.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor_test.go, line 23 at r1 (raw file):

Previously, rafiss (Rafi Shamim) wrote…

can this file also include a benchmark test?

Yes, I'll add one based on the benchmarks that I've done back then: https://docs.google.com/spreadsheets/d/1ui7Rkw9zIsEctOIFT8GFoO9aHpvBCJbBlsxg6JN_CmE/edit#gid=517332031.

[jaylim-crl]

TFTRs! The PR is ready for another round of review.

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @andy-kimball, @jeffswenson, and @rafiss)

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 56 at r1 (raw file):

Previously, jaylim-crl (Jay Lim) wrote…

The interceptor does not return any recoverable errors unless the caller wrapped src or dst to return custom errors. I've thought about this further, and I think I'll remove this. We'll just defer all of this logic to the caller. If the caller wraps any of those connections with custom errors, those errors should be handled accordingly. If this becomes an issue, we'll add this back.

I'm going to take that back. I'd rather be on the safe side, and keep closed here for correctness purposes. Once we're sure that everything has been stabilized, we can remove this. Similarly, if we decide in the future that this adds a lot of complexity to the caller, we can remove this.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 116 at r1 (raw file):

Previously, jaylim-crl (Jay Lim) wrote…

I'm not sure if it's worth handling this case given that we're already doing this on the DB side today:

cockroach/pkg/sql/pgwire/pgwirebase/encoding.go

Lines 154 to 156 in f1fcb83

	size := int(binary.BigEndian.Uint32(b.tmp[:]))
	// size includes itself.
	size -= 4

I'm thinking to leave this as-is unless you strongly feel about this.

As discussed, we'll leave it as it (i.e. cast int on uint32) given that we're already doing that elsewhere. That said, I've moved the subtraction to below, and checked on size < 4 instead.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 157 at r1 (raw file):

Previously, jaylim-crl (Jay Lim) wrote…

Good point. I can do that.

Done.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor_test.go, line 23 at r1 (raw file):

Previously, jaylim-crl (Jay Lim) wrote…

Yes, I'll add one based on the benchmarks that I've done back then: https://docs.google.com/spreadsheets/d/1ui7Rkw9zIsEctOIFT8GFoO9aHpvBCJbBlsxg6JN_CmE/edit#gid=517332031.

I added that in a different PR: #76429.

[andy-kimball]

just some nits / questions

Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @jaylim-crl, @jeffswenson, and @rafiss)

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 48 at r2 (raw file):

	// buf stores bytes which have been read, but have not been processed yet.
	buf []byte
	// readPos and writePos indicates the read and write pointers for bytes in

It'd be good to document a bit more detail on this buffer and pointers so that it's not necessary to read the code to discover the answers:

1. Does a message always start at the beginning of a buffer (i.e. 1st 5 bytes are always the message header)? Or is it a circular buffer?
2. Can the buffer contain multiple messages?
3. Is readPos always <= writePos?

Also, the comment says, "which have been read, but have not been processed yet", but if I call ReadMsg, that does "process" bytes in the buffer, correct?

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 75 at r2 (raw file):

// ensureNextNBytes blocks on IO reads until the buffer has at least n bytes.
func (p *pgInterceptor) ensureNextNBytes(n int) error {

NIT: I typically order exported methods above un-exported methods, roughly ordering from most important at the top to least important at the bottom, since humans typically read in that order (top-to-bottom).

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 168 at r2 (raw file):

		return nil, err
	}
	msgSizeBytes := pgHeaderSizeBytes + size

NIT: It is a bit strange that you're returning the size of the message minus the header, and then immediately adding back the header size. You're even calling the combined size msgSizeBytes. Maybe it'd be better to make PeekMsg return the unmodified message size to begin with?

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 188 at r2 (raw file):

	// Copy bytes which have already been read.
	toCopy := msgSizeBytes
	if p.ReadSize() < msgSizeBytes {

NIT: This condition will always be true, since msgSizeBytes > len(p.buf), and there's no way that p.ReadSize() could be >= that value.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 340 at r2 (raw file):

}

// FrontendInterceptor is a client interceptor for the Postgres frontend

NIT: Why not break one or more of these other classes into separate files?

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 343 at r2 (raw file):

// protocol.
type FrontendInterceptor struct {
	p *pgInterceptor

NIT: Should you instead use type FrontendInterceptor pgInterceptor here, to avoid allocating an extra object? Or embed the pgInterceptor struct rather than the pointer? Is there the possibility of adding more fields to this struct, or some other consideration? Not a big deal, though, since we won't create many of these.

[jaylim-crl]

TFTRs!

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @andy-kimball, @jeffswenson, and @rafiss)

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 48 at r2 (raw file):

1. No, a message can start in the middle of the buffer. This would happen if multiple messages were read into the buffer at once through the io.ReadAtLeast call. It is not a circular buffer for simplicity - we'll just read/write until the end of the buffer. All the public APIs handle this case.
2. Yes.
3. Yes, that is correct. If readPos == writePos, we have nothing in the buffer. writePos advances whenever we read into the buffer, whereas readPos advances whenever we read from the buffer to process the message; the latter is only possible through ReadMsg and ForwardMsg.

Also, the comment says, "which have been read, but have not been processed yet", but if I call ReadMsg, that does "process" bytes in the buffer, correct?

Yes, that is also correct.

I've updated the comment with more information.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 75 at r2 (raw file):

Previously, andy-kimball (Andy Kimball) wrote…

NIT: I typically order exported methods above un-exported methods, roughly ordering from most important at the top to least important at the bottom, since humans typically read in that order (top-to-bottom).

Done. Makes sense.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 168 at r2 (raw file):

Previously, andy-kimball (Andy Kimball) wrote…

NIT: It is a bit strange that you're returning the size of the message minus the header, and then immediately adding back the header size. You're even calling the combined size msgSizeBytes. Maybe it'd be better to make PeekMsg return the unmodified message size to begin with?

There are two types of sizes here:

1. message size: 1 byte (type) + 4 bytes (length) + # of body bytes
2. body size: # of body bytes

It's strange that the Postgres protocol defined the length int as the # of body bytes + 4 bytes (which is inclusive of itself). PeekMsg returns the body size here; note that we subtracted 4 bytes (since it included itself), and not the header size (5 bytes). If we returned the unmodified size (which is defined as 4 bytes + # of body bytes), we'd still have to add 1 for the type, and define that as message size here.

For now, I'll keep it as-is, i.e. PeekMsg returns the body size. I'll think further about this, and if we decide to change the definition of PeekMsg, I'll do that in a follow up PR. In that case, we'd still have to add 1 to the unmodified length for PeekMsg to return the message size.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 188 at r2 (raw file):

Previously, andy-kimball (Andy Kimball) wrote…

NIT: This condition will always be true, since msgSizeBytes > len(p.buf), and there's no way that p.ReadSize() could be >= that value.

Done. Good catch. In the initial version, we always allocate regardless, but now that we only allocate if the message doesn't fit, this condition is redundant.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 340 at r2 (raw file):

Previously, andy-kimball (Andy Kimball) wrote…

NIT: Why not break one or more of these other classes into separate files?

Done.

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 343 at r2 (raw file):

Previously, andy-kimball (Andy Kimball) wrote…

NIT: Should you instead use type FrontendInterceptor pgInterceptor here, to avoid allocating an extra object? Or embed the pgInterceptor struct rather than the pointer? Is there the possibility of adding more fields to this struct, or some other consideration? Not a big deal, though, since we won't create many of these.

No, we won't be adding more fields to this struct. The APIs that we have here should suffice. I moved it to use pgInterceptor as you suggested.

[jeffswenson]

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @andy-kimball, @jaylim-crl, and @rafiss)

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 168 at r2 (raw file):

Previously, jaylim-crl (Jay Lim) wrote…

There are two types of sizes here:

1. message size: 1 byte (type) + 4 bytes (length) + # of body bytes
2. body size: # of body bytes

It's strange that the Postgres protocol defined the length int as the # of body bytes + 4 bytes (which is inclusive of itself). PeekMsg returns the body size here; note that we subtracted 4 bytes (since it included itself), and not the header size (5 bytes). If we returned the unmodified size (which is defined as 4 bytes + # of body bytes), we'd still have to add 1 for the type, and define that as message size here.

For now, I'll keep it as-is, i.e. PeekMsg returns the body size. I'll think further about this, and if we decide to change the definition of PeekMsg, I'll do that in a follow up PR. In that case, we'd still have to add 1 to the unmodified length for PeekMsg to return the message size.

I'm okay with deferring this decision to a future PR, but I think the size returned should be bodySize + headerSize. All of the consumers of PeekMsg in this PR immediately add the header size. It also provides a nice symmetry in the API. That way the size returned by PeekMsg is equal to the length of the byte slice returned by ReadMsg().

[jaylim-crl]

TFTRs!

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @andy-kimball, @jeffswenson, and @rafiss)

---

pkg/ccl/sqlproxyccl/interceptor/interceptor.go, line 168 at r2 (raw file):

Previously, JeffSwenson (Jeff Swenson) wrote…

I'm okay with deferring this decision to a future PR, but I think the size returned should be bodySize + headerSize. All of the consumers of PeekMsg in this PR immediately add the header size. It also provides a nice symmetry in the API. That way the size returned by PeekMsg is equal to the length of the byte slice returned by ReadMsg().

Sounds like a plan. I added that as an item to #76000.

[jaylim-crl]

bors r=andy-kimball,JeffSwenson

[craig]

Build succeeded:

• GitHub CI (Cockroach)