go-libp2p-resource-manager/limit.go

package rcmgr

import (
	"encoding/json"
	"io"

	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/protocol"
)

// Limit is an object that specifies basic resource limits.
type Limit interface {
	// GetMemoryLimit returns the (current) memory limit.
	GetMemoryLimit() int64
	// GetStreamLimit returns the stream limit, for inbound or outbound streams.
	GetStreamLimit(network.Direction) int
	// GetStreamTotalLimit returns the total stream limit
	GetStreamTotalLimit() int
	// GetConnLimit returns the connection limit, for inbound or outbound connections.
	GetConnLimit(network.Direction) int
	// GetConnTotalLimit returns the total connection limit
	GetConnTotalLimit() int
	// GetFDLimit returns the file descriptor limit.
	GetFDLimit() int
}

// Limiter is the interface for providing limits to the resource manager.
type Limiter interface {
	GetSystemLimits() Limit
	GetTransientLimits() Limit
	GetServiceLimits(svc string) Limit
	GetServicePeerLimits(svc string) Limit
	GetProtocolLimits(proto protocol.ID) Limit
	GetProtocolPeerLimits(proto protocol.ID) Limit
	GetPeerLimits(p peer.ID) Limit
	GetStreamLimits(p peer.ID) Limit
	GetConnLimits() Limit
}

// NewDefaultLimiterFromJSON creates a new limiter by parsing a json configuration,
// using the default limits for fallback.
func NewDefaultLimiterFromJSON(in io.Reader) (Limiter, error) {
	return NewLimiterFromJSON(in, DefaultLimits.AutoScale())
}

// NewLimiterFromJSON creates a new limiter by parsing a json configuration.
func NewLimiterFromJSON(in io.Reader, defaults LimitConfig) (Limiter, error) {
	cfg, err := readLimiterConfigFromJSON(in, defaults)
	if err != nil {
		return nil, err
	}
	return &fixedLimiter{cfg}, nil
}

func readLimiterConfigFromJSON(in io.Reader, defaults LimitConfig) (LimitConfig, error) {
	var cfg LimitConfig
	if err := json.NewDecoder(in).Decode(&cfg); err != nil {
		return LimitConfig{}, err
	}
	cfg.Apply(defaults)
	return cfg, nil
}

// fixedLimiter is a limiter with fixed limits.
type fixedLimiter struct {
	LimitConfig
}

var _ Limiter = (*fixedLimiter)(nil)

func NewFixedLimiter(conf LimitConfig) Limiter {
	log.Debugw("initializing new limiter with config", "limits", conf)
	return &fixedLimiter{LimitConfig: conf}
}

// BaseLimit is a mixin type for basic resource limits.
type BaseLimit struct {
	Streams         int
	StreamsInbound  int
	StreamsOutbound int
	Conns           int
	ConnsInbound    int
	ConnsOutbound   int
	FD              int
	Memory          int64
}

// Apply overwrites all zero-valued limits with the values of l2
// Must not use a pointer receiver.
func (l *BaseLimit) Apply(l2 BaseLimit) {
	if l.Streams == 0 {
		l.Streams = l2.Streams
	}
	if l.StreamsInbound == 0 {
		l.StreamsInbound = l2.StreamsInbound
	}
	if l.StreamsOutbound == 0 {
		l.StreamsOutbound = l2.StreamsOutbound
	}
	if l.Conns == 0 {
		l.Conns = l2.Conns
	}
	if l.ConnsInbound == 0 {
		l.ConnsInbound = l2.ConnsInbound
	}
	if l.ConnsOutbound == 0 {
		l.ConnsOutbound = l2.ConnsOutbound
	}
	if l.Memory == 0 {
		l.Memory = l2.Memory
	}
	if l.FD == 0 {
		l.FD = l2.FD
	}
}

// BaseLimitIncrease is the increase per GB of system memory.
type BaseLimitIncrease struct {
	Streams         int
	StreamsInbound  int
	StreamsOutbound int
	Conns           int
	ConnsInbound    int
	ConnsOutbound   int
	Memory          int64
	FDFraction      float64
}

// Apply overwrites all zero-valued limits with the values of l2
// Must not use a pointer receiver.
func (l *BaseLimitIncrease) Apply(l2 BaseLimitIncrease) {
	if l.Streams == 0 {
		l.Streams = l2.Streams
	}
	if l.StreamsInbound == 0 {
		l.StreamsInbound = l2.StreamsInbound
	}
	if l.StreamsOutbound == 0 {
		l.StreamsOutbound = l2.StreamsOutbound
	}
	if l.Conns == 0 {
		l.Conns = l2.Conns
	}
	if l.ConnsInbound == 0 {
		l.ConnsInbound = l2.ConnsInbound
	}
	if l.ConnsOutbound == 0 {
		l.ConnsOutbound = l2.ConnsOutbound
	}
	if l.Memory == 0 {
		l.Memory = l2.Memory
	}
	if l.FDFraction == 0 {
		l.FDFraction = l2.FDFraction
	}
}

func (l *BaseLimit) GetStreamLimit(dir network.Direction) int {
	if dir == network.DirInbound {
		return l.StreamsInbound
	} else {
		return l.StreamsOutbound
	}
}

func (l *BaseLimit) GetStreamTotalLimit() int {
	return l.Streams
}

func (l *BaseLimit) GetConnLimit(dir network.Direction) int {
	if dir == network.DirInbound {
		return l.ConnsInbound
	} else {
		return l.ConnsOutbound
	}
}

func (l *BaseLimit) GetConnTotalLimit() int {
	return l.Conns
}

func (l *BaseLimit) GetFDLimit() int {
	return l.FD
}

func (l *BaseLimit) GetMemoryLimit() int64 {
	return l.Memory
}

func (l *fixedLimiter) GetSystemLimits() Limit {
	return &l.System
}

func (l *fixedLimiter) GetTransientLimits() Limit {
	return &l.Transient
}

func (l *fixedLimiter) GetServiceLimits(svc string) Limit {
	sl, ok := l.Service[svc]
	if !ok {
		return &l.ServiceDefault
	}
	return &sl
}

func (l *fixedLimiter) GetServicePeerLimits(svc string) Limit {
	pl, ok := l.ServicePeer[svc]
	if !ok {
		return &l.ServicePeerDefault
	}
	return &pl
}

func (l *fixedLimiter) GetProtocolLimits(proto protocol.ID) Limit {
	pl, ok := l.Protocol[proto]
	if !ok {
		return &l.ProtocolDefault
	}
	return &pl
}

func (l *fixedLimiter) GetProtocolPeerLimits(proto protocol.ID) Limit {
	pl, ok := l.ProtocolPeer[proto]
	if !ok {
		return &l.ProtocolPeerDefault
	}
	return &pl
}

func (l *fixedLimiter) GetPeerLimits(p peer.ID) Limit {
	pl, ok := l.Peer[p]
	if !ok {
		return &l.PeerDefault
	}
	return &pl
}

func (l *fixedLimiter) GetStreamLimits(_ peer.ID) Limit {
	return &l.Stream
}

func (l *fixedLimiter) GetConnLimits() Limit {
	return &l.Conn
}