wehub-resource-sync
136 lines
4.4 KiB
Go
136 lines
4.4 KiB
Go
// Copyright 2024 Dolthub, Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package main
|
|
|
|
import (
|
|
"errors"
|
|
"net"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5/pgproto3"
|
|
)
|
|
|
|
// Connection contains the connection that is used by the replay system for interacting with the Doltgres server.
|
|
type Connection struct {
|
|
frontend *pgproto3.Frontend
|
|
connection net.Conn
|
|
timeout time.Duration
|
|
reader *MessageReader
|
|
startup *pgproto3.StartupMessage
|
|
bmChan chan pgproto3.BackendMessage
|
|
errChan chan error
|
|
}
|
|
|
|
// NewConnection returns a new Connection.
|
|
func NewConnection(network string, reader *MessageReader, timeout time.Duration) (*Connection, error) {
|
|
connection, err := (&net.Dialer{}).Dial("tcp", network)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &Connection{
|
|
frontend: pgproto3.NewFrontend(connection, connection),
|
|
connection: connection,
|
|
timeout: timeout,
|
|
reader: reader,
|
|
startup: nil,
|
|
bmChan: make(chan pgproto3.BackendMessage),
|
|
errChan: make(chan error),
|
|
}, nil
|
|
}
|
|
|
|
// Close closes the internal connection.
|
|
func (c *Connection) Close() {
|
|
_ = c.connection.Close()
|
|
}
|
|
|
|
// EmptyReceiveBuffer empties the buffer used by Receive. Returns an error if the buffer was not empty.
|
|
func (c *Connection) EmptyReceiveBuffer() error {
|
|
if c.frontend.ReadBufferLen() > 0 {
|
|
for c.frontend.ReadBufferLen() > 0 {
|
|
_, _ = c.frontend.Receive()
|
|
}
|
|
return errors.New("Doltgres sent additional messages after ReadyForQuery")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Queue adds the given messages to the queue, which will be sent when Send is called. This is generally useful for when
|
|
// messages may conditionally be added, and we want to send all messages at once rather than one at a time.
|
|
func (c *Connection) Queue(messages ...pgproto3.FrontendMessage) {
|
|
for _, message := range messages {
|
|
if startupMessage, ok := message.(*pgproto3.StartupMessage); ok {
|
|
c.startup = startupMessage
|
|
}
|
|
c.frontend.Send(message)
|
|
}
|
|
}
|
|
|
|
// Receive returns the next message from the backend. If an error is returned, then the connection has been closed, and
|
|
// a new one should be opened.
|
|
func (c *Connection) Receive() (pgproto3.BackendMessage, error) {
|
|
var message pgproto3.BackendMessage
|
|
go func() {
|
|
var err error
|
|
message, err = c.frontend.Receive()
|
|
c.errChan <- err
|
|
}()
|
|
return message, c.handleErrorChannel(false, false)
|
|
}
|
|
|
|
// Send sends the given messages (or messages added via Queue) over the wire. Will sync the reader to the next query on
|
|
// error. If such behavior is not desired, then use SendNoSync. If an error is returned, then the connection has been
|
|
// closed, and a new one should be opened.
|
|
func (c *Connection) Send(messages ...pgproto3.FrontendMessage) error {
|
|
c.Queue(messages...)
|
|
go func() {
|
|
c.errChan <- c.frontend.Flush()
|
|
}()
|
|
return c.handleErrorChannel(true, true)
|
|
}
|
|
|
|
// SendNoSync is the same as Send, except that the reader will NOT sync to the next query if an error is encountered.
|
|
// This should only be called when we can guarantee that the reader is already synchronized to the next query. If an
|
|
// error is returned, then the connection has been closed, and a new one should be opened.
|
|
func (c *Connection) SendNoSync(messages ...pgproto3.FrontendMessage) error {
|
|
c.Queue(messages...)
|
|
go func() {
|
|
c.errChan <- c.frontend.Flush()
|
|
}()
|
|
return c.handleErrorChannel(false, true)
|
|
}
|
|
|
|
// handleErrorChannel handles whether the Connection should close once an item has been sent to the error channel. If an
|
|
// error is returned, then the Connection has been closed.
|
|
func (c *Connection) handleErrorChannel(syncOnError bool, isSend bool) error {
|
|
var err error
|
|
select {
|
|
case err = <-c.errChan:
|
|
case <-time.After(c.timeout):
|
|
if isSend {
|
|
err = errors.New("timeout during Send")
|
|
} else {
|
|
err = errors.New("timeout during Receive")
|
|
}
|
|
}
|
|
if err != nil {
|
|
if syncOnError {
|
|
c.reader.SyncToNextQuery()
|
|
}
|
|
c.reader.PushQueue(c.startup, &pgproto3.ReadyForQuery{TxStatus: 'I'})
|
|
_ = c.connection.Close()
|
|
}
|
|
return err
|
|
}
|