dragonfly/scheduler/supervisor/task.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 supervisor

import (
	"sync"
	"time"

	"d7y.io/dragonfly/v2/pkg/rpc/base"
	"d7y.io/dragonfly/v2/pkg/structure/sortedlist"
	"go.uber.org/atomic"
)

type TaskStatus uint8

func (status TaskStatus) String() string {
	switch status {
	case TaskStatusWaiting:
		return "Waiting"
	case TaskStatusRunning:
		return "Running"
	case TaskStatusSeeding:
		return "Seeding"
	case TaskStatusSuccess:
		return "Success"
	case TaskStatusZombie:
		return "Zombie"
	case TaskStatusFail:
		return "Fail"
	default:
		return "unknown"
	}
}

const (
	TaskStatusWaiting TaskStatus = iota
	TaskStatusRunning
	TaskStatusSeeding
	TaskStatusSuccess
	TaskStatusZombie
	TaskStatusFail
)

type Task struct {
	lock                 sync.RWMutex
	TaskID               string
	URL                  string
	URLMeta              *base.UrlMeta
	DirectPiece          []byte
	CreateTime           time.Time
	lastAccessTime       time.Time
	lastTriggerTime      time.Time
	pieceList            map[int32]*base.PieceInfo
	PieceTotal           int32
	ContentLength        int64
	status               TaskStatus
	peers                *sortedlist.SortedList
	backSourceLimit      atomic.Int32
	needClientBackSource atomic.Bool
	backSourcePeers      []string
	// TODO add cdnPeers
}

func NewTask(taskID, url string, meta *base.UrlMeta) *Task {
	return &Task{
		TaskID:     taskID,
		URL:        url,
		URLMeta:    meta,
		CreateTime: time.Now(),
		pieceList:  make(map[int32]*base.PieceInfo),
		peers:      sortedlist.NewSortedList(),
		status:     TaskStatusWaiting,
	}
}

func (task *Task) AddPeer(peer *Peer) {
	task.peers.UpdateOrAdd(peer)
}

func (task *Task) UpdatePeer(peer *Peer) {
	task.peers.Update(peer)
}

func (task *Task) DeletePeer(peer *Peer) {
	task.peers.Delete(peer)
}

func (task *Task) ListPeers() *sortedlist.SortedList {
	return task.peers
}

func (task *Task) SetStatus(status TaskStatus) {
	task.lock.Lock()
	defer task.lock.Unlock()
	task.status = status
}

func (task *Task) GetStatus() TaskStatus {
	task.lock.RLock()
	defer task.lock.RUnlock()
	return task.status
}

func (task *Task) SetClientBackSourceStatusAndLimit(backSourceLimit int32) {
	task.lock.Lock()
	defer task.lock.Unlock()
	task.backSourcePeers = make([]string, 0, backSourceLimit)
	task.needClientBackSource.Store(true)
	task.backSourceLimit.Store(backSourceLimit)
}

func (task *Task) NeedClientBackSource() bool {
	return task.needClientBackSource.Load()
}

func (task *Task) GetPiece(pieceNum int32) *base.PieceInfo {
	task.lock.RLock()
	defer task.lock.RUnlock()
	return task.pieceList[pieceNum]
}

func (task *Task) AddPiece(p *base.PieceInfo) {
	task.lock.RLock()
	if _, ok := task.pieceList[p.PieceNum]; ok {
		task.lock.RUnlock()
		return
	}
	task.lock.RUnlock()
	task.lock.Lock()
	defer task.lock.Unlock()
	task.pieceList[p.PieceNum] = p
}

func (task *Task) GetLastTriggerTime() time.Time {
	return task.lastTriggerTime
}

func (task *Task) UpdateLastTriggerTime(lastTriggerTime time.Time) {
	task.lastTriggerTime = lastTriggerTime
}

func (task *Task) Touch() {
	task.lock.Lock()
	defer task.lock.Unlock()
	task.lastAccessTime = time.Now()
}

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

func (task *Task) UpdateTaskSuccessResult(pieceTotal int32, contentLength int64) {
	task.lock.Lock()
	defer task.lock.Unlock()
	if task.status != TaskStatusSuccess {
		task.status = TaskStatusSuccess
		task.PieceTotal = pieceTotal
		task.ContentLength = contentLength
	}
}

func (task *Task) Lock() {
	task.lock.Lock()
}

func (task *Task) UnLock() {
	task.lock.Unlock()
}

func (task *Task) RLock() {
	task.lock.RLock()
}

func (task *Task) RUnlock() {
	task.lock.RUnlock()
}

const TinyFileSize = 128

// IsSuccess determines that whether cdn status is success.
func (task *Task) IsSuccess() bool {
	return task.status == TaskStatusSuccess
}

// IsFrozen determines that whether cdn status is frozen
//func (task *Task) IsFrozen() bool {
//	return task.status == TaskStatusWaiting || task.status == TaskStatusZombie || task.status == TaskStatusFailed ||
//		task.status == TaskStatusSourceError || task.status == TaskStatusCDNRegisterFail
//}

// CanSchedule determines whether task can be scheduled
// only task status is seeding or success can be scheduled
func (task *Task) CanSchedule() bool {
	return task.status == TaskStatusSeeding || task.status == TaskStatusSuccess
}

// IsWaiting determines whether task is waiting
func (task *Task) IsWaiting() bool {
	return task.status == TaskStatusWaiting
}

// IsHealth determines whether task is health
func (task *Task) IsHealth() bool {
	return task.status == TaskStatusRunning || task.status == TaskStatusSeeding || task.status == TaskStatusSuccess
}

// IsFail determines whether task is fail
func (task *Task) IsFail() bool {
	return task.status == TaskStatusFail
}

func (task *Task) IncreaseBackSourcePeer(peerID string) {
	task.lock.Lock()
	defer task.lock.Unlock()
	task.backSourcePeers = append(task.backSourcePeers, peerID)
	if task.backSourceLimit.Dec() <= 0 {
		task.needClientBackSource.Store(false)
	}
}

func (task *Task) GetBackSourcePeers() []string {
	task.lock.RLock()
	defer task.lock.RUnlock()
	backSourcePeers := task.backSourcePeers
	return backSourcePeers
}

func (task *Task) IsBackSourcePeer(peerID string) bool {
	task.lock.RLock()
	defer task.lock.RUnlock()
	for i := range task.backSourcePeers {
		if task.backSourcePeers[i] == peerID {
			return true
		}
	}
	return false
}

func (task *Task) Pick(limit int, pickFn func(peer *Peer) bool) (pickedPeers []*Peer) {
	return task.pick(limit, false, pickFn)
}

func (task *Task) PickReverse(limit int, pickFn func(peer *Peer) bool) (pickedPeers []*Peer) {
	return task.pick(limit, true, pickFn)
}

func (task *Task) pick(limit int, reverse bool, pickFn func(peer *Peer) bool) (pickedPeers []*Peer) {
	if pickFn == nil {
		return
	}
	if !reverse {
		task.ListPeers().Range(func(data sortedlist.Item) bool {
			if len(pickedPeers) >= limit {
				return false
			}
			peer := data.(*Peer)
			if pickFn(peer) {
				pickedPeers = append(pickedPeers, peer)
			}
			return true
		})
		return
	}
	task.ListPeers().RangeReverse(func(data sortedlist.Item) bool {
		if len(pickedPeers) >= limit {
			return false
		}
		peer := data.(*Peer)
		if pickFn(peer) {
			pickedPeers = append(pickedPeers, peer)
		}
		return true
	})
	return
}