dragonfly/cdn/supervisor/progress/progress.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 progress

import (
	"container/list"
	"context"
	"sort"
	"sync"

	"go.uber.org/atomic"

	"d7y.io/dragonfly/v2/cdn/supervisor/task"
	logger "d7y.io/dragonfly/v2/internal/dflog"
)

type subscriber struct {
	ctx       context.Context
	scheduler string
	taskID    string
	done      chan struct{}
	once      sync.Once
	pieces    map[uint32]*task.PieceInfo
	pieceChan chan *task.PieceInfo
	cond      *sync.Cond
	closed    *atomic.Bool
}

func newProgressSubscriber(ctx context.Context, clientAddr, taskID string, taskPieces map[uint32]*task.PieceInfo) *subscriber {
	pieces := make(map[uint32]*task.PieceInfo, len(taskPieces))
	for u, info := range taskPieces {
		pieces[u] = info
	}
	sub := &subscriber{
		ctx:       ctx,
		scheduler: clientAddr,
		taskID:    taskID,
		pieces:    pieces,
		done:      make(chan struct{}),
		pieceChan: make(chan *task.PieceInfo, 100),
		cond:      sync.NewCond(&sync.Mutex{}),
		closed:    atomic.NewBool(false),
	}
	go sub.readLoop()
	return sub
}

func (sub *subscriber) readLoop() {
	logger.Debugf("subscriber %s starts watching task %s seed progress", sub.scheduler, sub.taskID)
	defer func() {
		close(sub.pieceChan)
		logger.Debugf("subscriber %s stopped watch task %s seed progress", sub.scheduler, sub.taskID)
	}()
	for {
		select {
		case <-sub.ctx.Done():
			return
		case <-sub.done:
			if len(sub.pieces) == 0 {
				return
			}
			logger.Debugf("sub has been closed, there are still has %d pieces waiting to be sent", len(sub.pieces))
			sub.cond.L.Lock()
			sub.sendPieces()
			sub.cond.L.Unlock()
		default:
			sub.cond.L.Lock()
			for len(sub.pieces) == 0 && !sub.closed.Load() {
				sub.cond.Wait()
			}
			sub.sendPieces()
			sub.cond.L.Unlock()
		}
	}
}

func (sub *subscriber) sendPieces() {
	pieceNums := make([]uint32, 0, len(sub.pieces))
	for pieceNum := range sub.pieces {
		pieceNums = append(pieceNums, pieceNum)
	}
	sort.Slice(pieceNums, func(i, j int) bool {
		return pieceNums[i] < pieceNums[j]
	})
	for _, pieceNum := range pieceNums {
		logger.Debugf("subscriber %s send %d piece info of taskID %s", sub.scheduler, pieceNum, sub.taskID)
		sub.pieceChan <- sub.pieces[pieceNum]
		delete(sub.pieces, pieceNum)
	}
}

func (sub *subscriber) Notify(seedPiece *task.PieceInfo) {
	logger.Debugf("notifies subscriber %s about %d piece info of taskID %s", sub.scheduler, seedPiece.PieceNum, sub.taskID)
	sub.cond.L.Lock()
	sub.pieces[seedPiece.PieceNum] = seedPiece
	sub.cond.L.Unlock()
	sub.cond.Signal()
}

func (sub *subscriber) Receiver() <-chan *task.PieceInfo {
	return sub.pieceChan
}

func (sub *subscriber) Close() {
	sub.once.Do(func() {
		logger.Debugf("close subscriber %s from taskID %s", sub.scheduler, sub.taskID)
		sub.closed.CAS(false, true)
		sub.cond.Signal()
		close(sub.done)
	})
}

type publisher struct {
	taskID      string
	subscribers *list.List
}

func newProgressPublisher(taskID string) *publisher {
	return &publisher{
		taskID:      taskID,
		subscribers: list.New(),
	}
}

func (pub *publisher) AddSubscriber(sub *subscriber) {
	pub.subscribers.PushBack(sub)
	logger.Debugf("subscriber %s has been added into subscribers of publisher %s, list size is %d", sub.scheduler, sub.taskID, pub.subscribers.Len())
}

func (pub *publisher) RemoveSubscriber(sub *subscriber) {
	sub.Close()
	for e := pub.subscribers.Front(); e != nil; e = e.Next() {
		if e.Value == sub {
			pub.subscribers.Remove(e)
			logger.Debugf("subscriber %s has been removed from subscribers of publisher %s, list size is %d", sub.scheduler, sub.taskID, pub.subscribers.Len())
			return
		}
	}
}

func (pub *publisher) NotifySubscribers(seedPiece *task.PieceInfo) {
	for e := pub.subscribers.Front(); e != nil; e = e.Next() {
		e.Value.(*subscriber).Notify(seedPiece)
	}
}

func (pub *publisher) RemoveAllSubscribers() {
	var next *list.Element
	for e := pub.subscribers.Front(); e != nil; e = next {
		next = e.Next()
		pub.RemoveSubscriber(e.Value.(*subscriber))
	}
}