config: add a read limit option to the connection.
By default, the read limit is never set by NewConnection. When using the default pool dialer, a read limit of 512KB is set. When the read limit option is set, all sockets passed to NewConnection have that read limit set.
This commit is contained in:
+21
-1
@@ -51,6 +51,7 @@ type ConnectionConfig struct {
|
||||
PingInterval time.Duration
|
||||
IncomingBufferSize int
|
||||
ErrorsBufferSize int
|
||||
ReadLimit *int64
|
||||
}
|
||||
|
||||
func NewConnectionConfig(opts ...ConnectionOption) (*ConnectionConfig, error) {
|
||||
@@ -59,6 +60,7 @@ func NewConnectionConfig(opts ...ConnectionOption) (*ConnectionConfig, error) {
|
||||
PingInterval: 20 * time.Second,
|
||||
IncomingBufferSize: 100,
|
||||
ErrorsBufferSize: 10,
|
||||
ReadLimit: nil,
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(cfg)
|
||||
@@ -97,6 +99,14 @@ func WithErrorsBufferSize(value int) ConnectionOption {
|
||||
}
|
||||
}
|
||||
|
||||
// When ReadLimit is set to zero, reads are unlimited.
|
||||
func WithReadLimit(value int64) ConnectionOption {
|
||||
return func(c *ConnectionConfig) {
|
||||
ptr := &value
|
||||
c.ReadLimit = ptr
|
||||
}
|
||||
}
|
||||
|
||||
func ValidateConnectionConfig(c ConnectionConfig) error {
|
||||
if c.WriteTimeout < 0 {
|
||||
return fmt.Errorf("invalid write timeout: %v", c.WriteTimeout)
|
||||
@@ -112,6 +122,9 @@ func ValidateConnectionConfig(c ConnectionConfig) error {
|
||||
return fmt.Errorf("invalid errors buffer size: %d",
|
||||
c.ErrorsBufferSize)
|
||||
}
|
||||
if c.ReadLimit != nil && *c.ReadLimit < 0 {
|
||||
return fmt.Errorf("invalid read limit: %d", *c.ReadLimit)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -146,7 +159,10 @@ type Connection struct {
|
||||
}
|
||||
|
||||
func NewConnection(
|
||||
ctx context.Context, socket types.Socket, config *ConnectionConfig, handler slog.Handler,
|
||||
ctx context.Context,
|
||||
socket types.Socket,
|
||||
config *ConnectionConfig,
|
||||
handler slog.Handler,
|
||||
) (*Connection, error) {
|
||||
if socket == nil {
|
||||
return nil, ErrNilSocket
|
||||
@@ -179,6 +195,10 @@ func NewConnection(
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
|
||||
if config.ReadLimit != nil {
|
||||
conn.socket.SetReadLimit(*config.ReadLimit)
|
||||
}
|
||||
|
||||
if handler != nil {
|
||||
comp := component.FromContext(ctx)
|
||||
conn.logger = slog.New(handler).With(slog.Any("component", comp))
|
||||
|
||||
Reference in New Issue
Block a user