131 lines
3.4 KiB
Go
131 lines
3.4 KiB
Go
/*
|
|
* Copyright 2024 The Dragonfly Authors
|
|
*
|
|
* 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 pex
|
|
|
|
import (
|
|
"sync"
|
|
|
|
dfdaemonv1 "d7y.io/api/v2/pkg/apis/dfdaemon/v1"
|
|
|
|
logger "d7y.io/dragonfly/v2/internal/dflog"
|
|
)
|
|
|
|
type peerPool struct {
|
|
lock *sync.RWMutex
|
|
// map[task id] map[host id] *DestPeer
|
|
tasks map[string]map[string]*DestPeer
|
|
}
|
|
|
|
func newPeerPool() *peerPool {
|
|
return &peerPool{
|
|
lock: &sync.RWMutex{},
|
|
tasks: map[string]map[string]*DestPeer{},
|
|
}
|
|
}
|
|
|
|
func (p *peerPool) Sync(nodeMeta *MemberMeta, peerExchangeData *dfdaemonv1.PeerExchangeData) {
|
|
p.lock.Lock()
|
|
defer p.lock.Unlock()
|
|
|
|
for _, peer := range peerExchangeData.PeerMetadatas {
|
|
p.sync(nodeMeta, peer)
|
|
}
|
|
}
|
|
|
|
func (p *peerPool) sync(nodeMeta *MemberMeta, peer *dfdaemonv1.PeerMetadata) {
|
|
peers, ok := p.tasks[peer.TaskId]
|
|
if !ok {
|
|
peers = map[string]*DestPeer{}
|
|
p.tasks[peer.TaskId] = peers
|
|
}
|
|
|
|
clean := func() {
|
|
delete(peers, peer.PeerId)
|
|
// clean empty task map
|
|
if len(peers) == 0 {
|
|
delete(p.tasks, peer.TaskId)
|
|
}
|
|
}
|
|
|
|
switch peer.State {
|
|
case dfdaemonv1.PeerState_Unknown:
|
|
logger.Warnf("receive unknown state peer %s/%s from %s/%s, is local: %t",
|
|
peer.TaskId, peer.PeerId, nodeMeta.IP, nodeMeta.HostID, nodeMeta.isLocal)
|
|
return
|
|
case dfdaemonv1.PeerState_Running:
|
|
peers[nodeMeta.HostID] = &DestPeer{
|
|
MemberMeta: nodeMeta,
|
|
PeerID: peer.PeerId,
|
|
}
|
|
logger.Debugf("receive running peer %s/%s from %s/%s, is local: %t",
|
|
peer.TaskId, peer.PeerId, nodeMeta.IP, nodeMeta.HostID, nodeMeta.isLocal)
|
|
case dfdaemonv1.PeerState_Success:
|
|
peers[nodeMeta.HostID] = &DestPeer{
|
|
MemberMeta: nodeMeta,
|
|
PeerID: peer.PeerId,
|
|
}
|
|
logger.Debugf("receive successful peer %s/%s from %s/%s, is local: %t",
|
|
peer.TaskId, peer.PeerId, nodeMeta.IP, nodeMeta.HostID, nodeMeta.isLocal)
|
|
case dfdaemonv1.PeerState_Failed, dfdaemonv1.PeerState_Deleted:
|
|
clean()
|
|
logger.Debugf("receive deleted peer %s/%s from %s/%s, is local: %t",
|
|
peer.TaskId, peer.PeerId, nodeMeta.IP, nodeMeta.HostID, nodeMeta.isLocal)
|
|
default:
|
|
logger.Warnf("receive unknown state peer %s/%s from %s/%s, is local: %t",
|
|
peer.TaskId, peer.PeerId, nodeMeta.IP, nodeMeta.HostID, nodeMeta.isLocal)
|
|
return
|
|
}
|
|
}
|
|
|
|
func (p *peerPool) Search(task string) SearchPeerResult {
|
|
p.lock.RLock()
|
|
defer p.lock.RUnlock()
|
|
peers, ok := p.tasks[task]
|
|
if !ok || len(peers) == 0 {
|
|
return SearchPeerResult{
|
|
Type: SearchPeerResultTypeNotFound,
|
|
}
|
|
}
|
|
|
|
var (
|
|
dp []*DestPeer
|
|
typ SearchPeerResultType = SearchPeerResultTypeRemote
|
|
)
|
|
for _, peer := range peers {
|
|
// put local peer first
|
|
if peer.isLocal {
|
|
typ = SearchPeerResultTypeLocal
|
|
dp = append([]*DestPeer{peer}, dp...)
|
|
} else {
|
|
dp = append(dp, peer)
|
|
}
|
|
}
|
|
|
|
return SearchPeerResult{
|
|
Type: typ,
|
|
Peers: dp,
|
|
}
|
|
}
|
|
|
|
func (p *peerPool) clean(hostID string) {
|
|
for _, peers := range p.tasks {
|
|
if _, ok := peers[hostID]; ok {
|
|
delete(peers, hostID)
|
|
}
|
|
}
|
|
}
|