dragonfly/scheduler/types/peer.go

/*
 *     Copyright 2020 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 types

import (
	"sync"
	"time"

	"d7y.io/dragonfly/v2/pkg/rpc/scheduler"
	"go.uber.org/atomic"
)

type PeerStatus uint8

func (status PeerStatus) String() string {
	switch status {
	case PeerStatusWaiting:
		return "Waiting"
	case PeerStatusRunning:
		return "Running"
	case PeerStatusSuccess:
		return "Success"
	case PeerStatusFail:
		return "fail"
	case PeerStatusZombie:
		return "zombie"
	default:
		return "unknown"
	}
}

const (
	PeerStatusWaiting PeerStatus = iota
	PeerStatusRunning
	PeerStatusZombie
	PeerStatusFail
	PeerStatusSuccess
)

type Peer struct {
	lock sync.RWMutex
	// PeerID specifies ID of peer
	PeerID string
	// Task specifies
	Task *Task
	// Host specifies
	Host *PeerHost
	// bindPacketChan
	bindPacketChan bool
	// PacketChan send schedulerPacket to peer client
	packetChan chan *scheduler.PeerPacket
	// createTime
	CreateTime time.Time
	// finishedNum specifies downloaded finished piece number
	finishedNum    atomic.Int32
	lastAccessTime time.Time
	parent         *Peer
	children       sync.Map
	status         PeerStatus
	costHistory    []int
	leave          atomic.Bool
}

func NewPeer(peerID string, task *Task, host *PeerHost) *Peer {
	return &Peer{
		PeerID:         peerID,
		Task:           task,
		Host:           host,
		CreateTime:     time.Now(),
		lastAccessTime: time.Now(),
		status:         PeerStatusWaiting,
	}
}

func (peer *Peer) GetWholeTreeNode() int {
	// TODO lock task
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	count := 1
	peer.children.Range(func(key, value interface{}) bool {
		peerNode := value.(*Peer)
		count += peerNode.GetWholeTreeNode()
		return true
	})
	return count
}

func (peer *Peer) GetLastAccessTime() time.Time {
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	return peer.lastAccessTime
}

func (peer *Peer) Touch() {
	peer.lock.Lock()
	defer peer.lock.Unlock()
	peer.lastAccessTime = time.Now()
	if peer.status == PeerStatusZombie && !peer.leave.Load() {
		peer.status = PeerStatusRunning
	}
	peer.Task.Touch()
}

func (peer *Peer) associateChild(child *Peer) {
	peer.lock.Lock()
	peer.children.Store(child.PeerID, child)
	peer.Host.IncUploadLoad()
	peer.lock.Unlock()
	peer.Task.peers.Update(peer)
}

func (peer *Peer) disassociateChild(child *Peer) {
	peer.lock.Lock()
	peer.children.Delete(child.PeerID)
	peer.Host.DecUploadLoad()
	peer.lock.Unlock()
	peer.Task.peers.Update(peer)
}

func (peer *Peer) ReplaceParent(parent *Peer) {
	oldParent := peer.parent
	if oldParent != nil {
		oldParent.disassociateChild(peer)
	}
	peer.parent = parent
	if parent != nil {
		parent.associateChild(peer)
	}
}

func (peer *Peer) GetCostHistory() []int {
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	return peer.costHistory
}

func (peer *Peer) GetCost() int {
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	if len(peer.costHistory) < 1 {
		return int(time.Second / time.Millisecond)
	}
	totalCost := 0
	for _, cost := range peer.costHistory {
		totalCost += cost
	}
	return totalCost / len(peer.costHistory)
}

func (peer *Peer) AddPieceInfo(finishedCount int32, cost int) {
	peer.lock.Lock()
	if finishedCount > peer.finishedNum.Load() {
		peer.finishedNum.Store(finishedCount)
		peer.costHistory = append(peer.costHistory, cost)
		if len(peer.costHistory) > 20 {
			peer.costHistory = peer.costHistory[len(peer.costHistory)-20:]
		}
		peer.lock.Unlock()
		peer.Task.peers.Update(peer)
		return
	}
	peer.lock.Unlock()
}

func (peer *Peer) GetDepth() int {
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	var deep int
	node := peer
	for node != nil {
		deep++
		if node.parent == nil || node.Host.CDN {
			break
		}
		node = node.parent
	}
	return deep
}

func (peer *Peer) GetTreeRoot() *Peer {
	node := peer
	for node != nil {
		if node.parent == nil || node.Host.CDN {
			break
		}
		node = node.parent
	}
	return node
}

// if peer is offspring of ancestor
func (peer *Peer) IsDescendantOf(ancestor *Peer) bool {
	if ancestor == nil {
		return false
	}
	// TODO avoid circulation
	peer.lock.RLock()
	ancestor.lock.RLock()
	defer ancestor.lock.RUnlock()
	defer peer.lock.RUnlock()
	node := peer
	for node != nil {
		if node.parent == nil || node.Host.CDN {
			return false
		} else if node.PeerID == ancestor.PeerID {
			return true
		}
		node = node.parent
	}
	return false
}

func (peer *Peer) IsAncestorOf(offspring *Peer) bool {
	if offspring == nil {
		return false
	}
	offspring.lock.RLock()
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	defer offspring.lock.RUnlock()
	node := offspring
	for node != nil {
		if node.parent == nil || node.Host.CDN {
			return false
		} else if node.PeerID == peer.PeerID {
			return true
		}
		node = node.parent
	}
	return false
}

func (peer *Peer) IsBlocking() bool {
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	if peer.parent == nil {
		return false
	}
	return peer.finishedNum.Load() >= peer.parent.finishedNum.Load()
}

func (peer *Peer) GetSortKeys() (key1, key2 int) {
	key1 = int(peer.finishedNum.Load())
	key2 = peer.getFreeLoad()
	return
}

func (peer *Peer) getFreeLoad() int {
	if peer.Host == nil {
		return 0
	}
	return peer.Host.GetFreeUploadLoad()
}

func GetDiffPieceNum(src *Peer, dst *Peer) int32 {
	diff := src.finishedNum.Load() - dst.finishedNum.Load()
	if diff > 0 {
		return diff
	}
	return -diff
}

func (peer *Peer) GetParent() *Peer {
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	return peer.parent
}

func (peer *Peer) GetChildren() *sync.Map {
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	return &peer.children
}

func (peer *Peer) SetStatus(status PeerStatus) {
	peer.lock.Lock()
	defer peer.lock.Unlock()
	peer.status = status
}

func (peer *Peer) BindSendChannel(packetChan chan *scheduler.PeerPacket) {
	peer.lock.Lock()
	defer peer.lock.Unlock()
	peer.bindPacketChan = true
	peer.packetChan = packetChan
}

func (peer *Peer) UnBindSendChannel() {
	peer.lock.Lock()
	defer peer.lock.Unlock()
	if peer.bindPacketChan {
		if peer.packetChan != nil {
			close(peer.packetChan)
		}
		peer.bindPacketChan = false
	}
}

func (peer *Peer) SendSchedulePacket(packet *scheduler.PeerPacket) {
	peer.lock.Lock()
	defer peer.lock.Unlock()
	if peer.bindPacketChan {
		peer.packetChan <- packet
	}
}

func (peer *Peer) IsRunning() bool {
	return peer.status == PeerStatusRunning
}

func (peer *Peer) IsWaiting() bool {
	return peer.status == PeerStatusWaiting
}

func (peer *Peer) IsSuccess() bool {
	return peer.status == PeerStatusSuccess
}

func (peer *Peer) IsDone() bool {
	return peer.status == PeerStatusSuccess || peer.status == PeerStatusFail
}

func (peer *Peer) IsBad() bool {
	return peer.status == PeerStatusFail || peer.status == PeerStatusZombie
}

func (peer *Peer) GetFinishedNum() int32 {
	return peer.finishedNum.Load()
}

func (peer *Peer) GetStatus() PeerStatus {
	peer.lock.RLock()
	defer peer.lock.RUnlock()
	return peer.status
}

func (peer *Peer) MarkLeave() {
	peer.leave.Store(true)
}

func (peer *Peer) IsLeave() bool {
	return peer.leave.Load()
}