<commit_before>package connmgr

import (
	"context"
	"sort"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

var log = logging.Logger("connmgr")

type ConnManager interface {
	TagPeer(peer.ID, string, int)
	UntagPeer(peer.ID, string)
	GetTagInfo(peer.ID) *TagInfo
	TrimOpenConns(context.Context)
	Notifee() inet.Notifiee
}

type connManager struct {
	highWater int
	lowWater  int

	gracePeriod time.Duration

	peers     map[peer.ID]*peerInfo
	connCount int

	lk sync.Mutex

	lastTrim time.Time
}

var DefaultGracePeriod = time.Second * 10

func NewConnManager(low, hi int, grace time.Duration) ConnManager {
	return &connManager{
		highWater:   hi,
		lowWater:    low,
		gracePeriod: grace,
		peers:       make(map[peer.ID]*peerInfo),
	}
}

type peerInfo struct {
	tags  map[string]int
	value int

	conns map[inet.Conn]time.Time

	firstSeen time.Time
}

type TagInfo struct {
	FirstSeen time.Time
	Value     int
	Tags      map[string]int
	Conns     map[string]time.Time
}

func (cm *connManager) TrimOpenConns(ctx context.Context) {
	cm.lk.Lock()
	defer cm.lk.Unlock()
	defer log.EventBegin(ctx, "connCleanup").Done()
	if cm.lowWater == 0 || cm.highWater == 0 {
		// disabled
		return
	}
	cm.lastTrim = time.Now()

	if len(cm.peers) < cm.lowWater {
		log.Info("open connection count below limit")
		return
	}

	var infos []*peerInfo

	for _, inf := range cm.peers {
		infos = append(infos, inf)
	}

	sort.Slice(infos, func(i, j int) bool {
		return infos[i].value < infos[j].value
	})

	close_count := len(infos) - cm.lowWater
	toclose := infos[:close_count]

	for _, inf := range toclose {
		if time.Since(inf.firstSeen) < cm.gracePeriod {
			continue
		}

		// TODO: if a peer has more than one connection, maybe only close one?
		for c, _ := range inf.conns {
			log.Info("closing conn: ", c.RemotePeer())
			log.Event(ctx, "closeConn", c.RemotePeer())
			c.Close()
		}
	}

	if len(cm.peers) > cm.highWater {
		log.Error("still over high water mark after trimming connections")
	}
}

func (cm *connManager) GetTagInfo(p peer.ID) *TagInfo {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		return nil
	}

	out := &TagInfo{
		FirstSeen: pi.firstSeen,
		Value:     pi.value,
		Tags:      make(map[string]int),
		Conns:     make(map[string]time.Time),
	}

	for t, v := range pi.tags {
		out.Tags[t] = v
	}
	for c, t := range pi.conns {
		out.Conns[c.RemoteMultiaddr().String()] = t
	}

	return out
}

func (cm *connManager) TagPeer(p peer.ID, tag string, val int) {<commit_msg>don't try to close a connection while holding a lock that blocks the notifiees<commit_after>package connmgr

import (
	"context"
	"io"
	"sort"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

var log = logging.Logger("connmgr")

type ConnManager interface {
	TagPeer(peer.ID, string, int)
	UntagPeer(peer.ID, string)
	GetTagInfo(peer.ID) *TagInfo
	TrimOpenConns(context.Context)
	Notifee() inet.Notifiee
}

type connManager struct {
	highWater int
	lowWater  int

	gracePeriod time.Duration

	peers     map[peer.ID]*peerInfo
	connCount int

	lk sync.Mutex

	lastTrim time.Time
}

var DefaultGracePeriod = time.Second * 10

func NewConnManager(low, hi int, grace time.Duration) ConnManager {
	return &connManager{
		highWater:   hi,
		lowWater:    low,
		gracePeriod: grace,
		peers:       make(map[peer.ID]*peerInfo),
	}
}

type peerInfo struct {
	tags  map[string]int
	value int

	conns map[inet.Conn]time.Time

	firstSeen time.Time
}

type TagInfo struct {
	FirstSeen time.Time
	Value     int
	Tags      map[string]int
	Conns     map[string]time.Time
}

func (cm *connManager) TrimOpenConns(ctx context.Context) {
	for _, c := range cm.getConnsToClose(ctx) {
		c.Close()
	}
}

func (cm *connManager) getConnsToClose(ctx context.Context) []io.Closer {
	cm.lk.Lock()
	defer cm.lk.Unlock()
	defer log.EventBegin(ctx, "connCleanup").Done()
	if cm.lowWater == 0 || cm.highWater == 0 {
		// disabled
		return nil
	}
	cm.lastTrim = time.Now()

	if len(cm.peers) < cm.lowWater {
		log.Info("open connection count below limit")
		return nil
	}

	var infos []*peerInfo

	for _, inf := range cm.peers {
		infos = append(infos, inf)
	}

	sort.Slice(infos, func(i, j int) bool {
		return infos[i].value < infos[j].value
	})

	close_count := len(infos) - cm.lowWater
	toclose := infos[:close_count]

	var closed []io.Closer

	for _, inf := range toclose {
		if time.Since(inf.firstSeen) < cm.gracePeriod {
			continue
		}

		// TODO: if a peer has more than one connection, maybe only close one?
		for c, _ := range inf.conns {
			log.Info("closing conn: ", c.RemotePeer())
			log.Event(ctx, "closeConn", c.RemotePeer())
			// TODO: probably don't want to always do this in a goroutine
			closed = append(closed, c)
		}
	}

	return closed
}

func (cm *connManager) GetTagInfo(p peer.ID) *TagInfo {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		return nil
	}

	out := &TagInfo{
		FirstSeen: pi.firstSeen,
		Value:     pi.value,
		Tags:      make(map[string]int),
		Conns:     make(map[string]time.Time),
	}

	for t, v := range pi.tags {
		out.Tags[t] = v
	}
	for c, t := range pi.conns {
		out.Conns[c.RemoteMultiaddr().String()] = t
	}

	return out
}

func (cm *connManager) TagPeer(p peer.ID, tag string, val int) {
{'fraction_non_alphanumeric': 0.13395297977036633, 'fraction_numerical': 0.00291598323309641, 'mean_word_length': 3.48, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '486927', 'n_tokens_mistral': 2258, 'n_tokens_neox': 2125, 'n_words': 508}