427 lines
12 KiB
Go
427 lines
12 KiB
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 peer
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"d7y.io/dragonfly/v2/pkg/dferrors"
|
|
"github.com/go-http-utils/headers"
|
|
"github.com/pkg/errors"
|
|
"go.opentelemetry.io/otel/semconv"
|
|
"go.opentelemetry.io/otel/trace"
|
|
"golang.org/x/time/rate"
|
|
|
|
"d7y.io/dragonfly/v2/client/config"
|
|
"d7y.io/dragonfly/v2/client/daemon/storage"
|
|
"d7y.io/dragonfly/v2/pkg/dfcodes"
|
|
logger "d7y.io/dragonfly/v2/pkg/dflog"
|
|
"d7y.io/dragonfly/v2/pkg/rpc/base"
|
|
"d7y.io/dragonfly/v2/pkg/rpc/scheduler"
|
|
schedulerclient "d7y.io/dragonfly/v2/pkg/rpc/scheduler/client"
|
|
)
|
|
|
|
// StreamPeerTask represents a peer task with stream io for reading directly without once more disk io
|
|
type StreamPeerTask interface {
|
|
Task
|
|
// Start start the special peer task, return a io.Reader for stream io
|
|
// when all data transferred, reader return a io.EOF
|
|
// attribute stands some extra data, like HTTP response Header
|
|
Start(ctx context.Context) (reader io.Reader, attribute map[string]string, err error)
|
|
}
|
|
|
|
type streamPeerTask struct {
|
|
peerTask
|
|
successPieceCh chan int32
|
|
}
|
|
|
|
func newStreamPeerTask(ctx context.Context,
|
|
host *scheduler.PeerHost,
|
|
pieceManager PieceManager,
|
|
request *scheduler.PeerTaskRequest,
|
|
schedulerClient schedulerclient.SchedulerClient,
|
|
schedulerOption config.SchedulerOption,
|
|
perPeerRateLimit rate.Limit) (context.Context, StreamPeerTask, *TinyData, error) {
|
|
ctx, span := tracer.Start(ctx, config.SpanStreamPeerTask, trace.WithSpanKind(trace.SpanKindClient))
|
|
span.SetAttributes(config.AttributePeerHost.String(host.Uuid))
|
|
span.SetAttributes(semconv.NetHostIPKey.String(host.Ip))
|
|
span.SetAttributes(config.AttributePeerID.String(request.PeerId))
|
|
span.SetAttributes(semconv.HTTPURLKey.String(request.Url))
|
|
|
|
logger.Debugf("request overview, url: %s, filter: %s, meta: %s, biz: %s", request.Url, request.Filter, request.UrlMata, request.BizId)
|
|
// trace register
|
|
_, regSpan := tracer.Start(ctx, config.SpanRegisterTask)
|
|
result, err := schedulerClient.RegisterPeerTask(ctx, request)
|
|
regSpan.RecordError(err)
|
|
regSpan.End()
|
|
|
|
var backSource bool
|
|
if err != nil {
|
|
// check if it is back source error
|
|
if de, ok := err.(*dferrors.DfError); ok && de.Code == dfcodes.SchedNeedBackSource {
|
|
backSource = true
|
|
}
|
|
// not back source
|
|
if !backSource {
|
|
logger.Errorf("register peer task failed: %s, peer id: %s", err, request.PeerId)
|
|
return ctx, nil, nil, err
|
|
}
|
|
}
|
|
if result == nil {
|
|
defer span.End()
|
|
span.RecordError(err)
|
|
err = errors.Errorf("empty schedule result")
|
|
return ctx, nil, nil, err
|
|
}
|
|
span.SetAttributes(config.AttributeTaskID.String(result.TaskId))
|
|
logger.Infof("register task success, task id: %s, peer id: %s, SizeScope: %s",
|
|
result.TaskId, request.PeerId, base.SizeScope_name[int32(result.SizeScope)])
|
|
|
|
var singlePiece *scheduler.SinglePiece
|
|
if !backSource {
|
|
switch result.SizeScope {
|
|
case base.SizeScope_SMALL:
|
|
span.SetAttributes(config.AttributePeerTaskSizeScope.String("small"))
|
|
logger.Debugf("%s/%s size scope: small", result.TaskId, request.PeerId)
|
|
if piece, ok := result.DirectPiece.(*scheduler.RegisterResult_SinglePiece); ok {
|
|
singlePiece = piece.SinglePiece
|
|
}
|
|
case base.SizeScope_TINY:
|
|
defer span.End()
|
|
span.SetAttributes(config.AttributePeerTaskSizeScope.String("tiny"))
|
|
logger.Debugf("%s/%s size scope: tiny", result.TaskId, request.PeerId)
|
|
if piece, ok := result.DirectPiece.(*scheduler.RegisterResult_PieceContent); ok {
|
|
return ctx, nil, &TinyData{
|
|
span: span,
|
|
TaskID: result.TaskId,
|
|
PeerID: request.PeerId,
|
|
Content: piece.PieceContent,
|
|
}, nil
|
|
}
|
|
err = errors.Errorf("scheduler return tiny piece but can not parse piece content")
|
|
span.RecordError(err)
|
|
return ctx, nil, nil, err
|
|
case base.SizeScope_NORMAL:
|
|
span.SetAttributes(config.AttributePeerTaskSizeScope.String("normal"))
|
|
logger.Debugf("%s/%s size scope: normal", result.TaskId, request.PeerId)
|
|
}
|
|
}
|
|
|
|
peerPacketStream, err := schedulerClient.ReportPieceResult(ctx, result.TaskId, request)
|
|
if err != nil {
|
|
defer span.End()
|
|
span.RecordError(err)
|
|
return ctx, nil, nil, err
|
|
}
|
|
var limiter *rate.Limiter
|
|
if perPeerRateLimit > 0 {
|
|
limiter = rate.NewLimiter(perPeerRateLimit, int(perPeerRateLimit))
|
|
}
|
|
return ctx, &streamPeerTask{
|
|
peerTask: peerTask{
|
|
ctx: ctx,
|
|
host: host,
|
|
backSource: backSource,
|
|
request: request,
|
|
peerPacketStream: peerPacketStream,
|
|
pieceManager: pieceManager,
|
|
peerPacketReady: make(chan bool, 1),
|
|
peerID: request.PeerId,
|
|
taskID: result.TaskId,
|
|
singlePiece: singlePiece,
|
|
done: make(chan struct{}),
|
|
span: span,
|
|
once: sync.Once{},
|
|
readyPieces: NewBitmap(),
|
|
requestedPieces: NewBitmap(),
|
|
lock: &sync.Mutex{},
|
|
failedPieceCh: make(chan int32, 4),
|
|
failedReason: "unknown",
|
|
failedCode: dfcodes.UnknownError,
|
|
contentLength: -1,
|
|
totalPiece: -1,
|
|
schedulerOption: schedulerOption,
|
|
limiter: limiter,
|
|
|
|
SugaredLoggerOnWith: logger.With("peer", request.PeerId, "task", result.TaskId, "component", "streamPeerTask"),
|
|
},
|
|
successPieceCh: make(chan int32, 4),
|
|
}, nil, nil
|
|
}
|
|
|
|
func (s *streamPeerTask) ReportPieceResult(piece *base.PieceInfo, pieceResult *scheduler.PieceResult) error {
|
|
defer s.recoverFromPanic()
|
|
// retry failed piece
|
|
if !pieceResult.Success {
|
|
_ = s.peerPacketStream.Send(pieceResult)
|
|
s.failedPieceCh <- pieceResult.PieceNum
|
|
return nil
|
|
}
|
|
|
|
s.lock.Lock()
|
|
if s.readyPieces.IsSet(pieceResult.PieceNum) {
|
|
s.lock.Unlock()
|
|
s.Warnf("piece %d is already reported, skipped", pieceResult.PieceNum)
|
|
return nil
|
|
}
|
|
// mark piece processed
|
|
s.readyPieces.Set(pieceResult.PieceNum)
|
|
atomic.AddInt64(&s.completedLength, int64(piece.RangeSize))
|
|
s.lock.Unlock()
|
|
|
|
pieceResult.FinishedCount = s.readyPieces.Settled()
|
|
_ = s.peerPacketStream.Send(pieceResult)
|
|
s.successPieceCh <- piece.PieceNum
|
|
s.Debugf("success piece %d sent", piece.PieceNum)
|
|
select {
|
|
case <-s.ctx.Done():
|
|
s.Warnf("peer task context done due to %s", s.ctx.Err())
|
|
return s.ctx.Err()
|
|
default:
|
|
}
|
|
|
|
if !s.isCompleted() {
|
|
return nil
|
|
}
|
|
|
|
return s.finish()
|
|
}
|
|
|
|
func (s *streamPeerTask) Start(ctx context.Context) (io.Reader, map[string]string, error) {
|
|
s.ctx, s.cancel = context.WithCancel(ctx)
|
|
if s.backSource {
|
|
go func() {
|
|
s.contentLength = -1
|
|
_ = s.callback.Init(s)
|
|
err := s.pieceManager.DownloadSource(ctx, s, s.request)
|
|
if err != nil {
|
|
s.Errorf("download from source error: %s", err)
|
|
s.cleanUnfinished()
|
|
return
|
|
}
|
|
s.Debugf("download from source ok")
|
|
_ = s.finish()
|
|
}()
|
|
} else {
|
|
s.pullPieces(s, s.cleanUnfinished)
|
|
}
|
|
|
|
// wait first piece to get content length and attribute (eg, response header for http/https)
|
|
var firstPiece int32
|
|
select {
|
|
case <-s.ctx.Done():
|
|
var err error
|
|
if s.failedReason != "" {
|
|
err = errors.Errorf(s.failedReason)
|
|
} else {
|
|
err = errors.Errorf("ctx.PeerTaskDone due to: %s", s.ctx.Err())
|
|
}
|
|
s.Errorf("%s", err)
|
|
s.span.RecordError(err)
|
|
s.span.End()
|
|
return nil, nil, err
|
|
case <-s.done:
|
|
var err error
|
|
if s.failedReason != "" {
|
|
err = errors.Errorf(s.failedReason)
|
|
} else {
|
|
err = errors.Errorf("stream peer task early done")
|
|
}
|
|
s.Errorf("%s", err)
|
|
s.span.RecordError(err)
|
|
s.span.End()
|
|
return nil, nil, err
|
|
case first := <-s.successPieceCh:
|
|
//if !ok {
|
|
// s.Warnf("successPieceCh closed unexpect")
|
|
// return nil, nil, errors.NewDaemonConfig("early done")
|
|
//}
|
|
firstPiece = first
|
|
}
|
|
|
|
pr, pw := io.Pipe()
|
|
attr := map[string]string{}
|
|
var reader io.Reader = pr
|
|
var writer io.Writer = pw
|
|
if s.contentLength != -1 {
|
|
attr[headers.ContentLength] = fmt.Sprintf("%d", s.contentLength)
|
|
} else {
|
|
attr[headers.TransferEncoding] = "chunked"
|
|
}
|
|
attr[config.HeaderDragonflyTask] = s.taskID
|
|
attr[config.HeaderDragonflyPeer] = s.peerID
|
|
|
|
go func(first int32) {
|
|
defer func() {
|
|
s.cancel()
|
|
s.span.End()
|
|
}()
|
|
var (
|
|
desired int32
|
|
cur int32
|
|
wrote int64
|
|
err error
|
|
//ok bool
|
|
cache = make(map[int32]bool)
|
|
)
|
|
// update first piece to cache and check cur with desired
|
|
cache[first] = true
|
|
cur = first
|
|
for {
|
|
if desired == cur {
|
|
for {
|
|
delete(cache, desired)
|
|
_, span := tracer.Start(s.ctx, config.SpanWriteBackPiece)
|
|
span.SetAttributes(config.AttributePiece.Int(int(desired)))
|
|
wrote, err = s.writeTo(writer, desired)
|
|
if err != nil {
|
|
span.RecordError(err)
|
|
span.End()
|
|
s.Errorf("write to pipe error: %s", err)
|
|
_ = pw.CloseWithError(err)
|
|
return
|
|
}
|
|
s.Debugf("wrote piece %d to pipe, size %d", desired, wrote)
|
|
span.End()
|
|
desired++
|
|
cached := cache[desired]
|
|
if !cached {
|
|
break
|
|
}
|
|
}
|
|
} else {
|
|
// not desired piece, cache it
|
|
cache[cur] = true
|
|
if cur < desired {
|
|
s.Warnf("piece number should be equal or greater than %d, received piece number: %d",
|
|
desired, cur)
|
|
}
|
|
}
|
|
|
|
select {
|
|
case <-s.ctx.Done():
|
|
s.Errorf("ctx.PeerTaskDone due to: %s", s.ctx.Err())
|
|
s.span.RecordError(s.ctx.Err())
|
|
if err := pw.CloseWithError(s.ctx.Err()); err != nil {
|
|
s.Errorf("CloseWithError failed: %s", err)
|
|
}
|
|
return
|
|
case <-s.done:
|
|
for {
|
|
// all data is wrote to local storage, and all data is wrote to pipe write
|
|
if s.readyPieces.Settled() == desired {
|
|
pw.Close()
|
|
return
|
|
}
|
|
_, span := tracer.Start(s.ctx, config.SpanWriteBackPiece)
|
|
span.SetAttributes(config.AttributePiece.Int(int(desired)))
|
|
wrote, err = s.writeTo(pw, desired)
|
|
if err != nil {
|
|
span.RecordError(err)
|
|
span.End()
|
|
s.span.RecordError(err)
|
|
s.Errorf("write to pipe error: %s", err)
|
|
_ = pw.CloseWithError(err)
|
|
return
|
|
}
|
|
span.End()
|
|
s.Debugf("wrote piece %d to pipe, size %d", desired, wrote)
|
|
desired++
|
|
}
|
|
case cur = <-s.successPieceCh:
|
|
continue
|
|
//if !ok {
|
|
// s.Warnf("successPieceCh closed")
|
|
// continue
|
|
//}
|
|
}
|
|
}
|
|
}(firstPiece)
|
|
|
|
return reader, attr, nil
|
|
}
|
|
|
|
func (s *streamPeerTask) finish() error {
|
|
// send last progress
|
|
s.once.Do(func() {
|
|
// send EOF piece result to scheduler
|
|
_ = s.peerPacketStream.Send(
|
|
scheduler.NewEndPieceResult(s.taskID, s.peerID, s.readyPieces.Settled()))
|
|
s.Debugf("end piece result sent")
|
|
close(s.done)
|
|
//close(s.successPieceCh)
|
|
if err := s.callback.Done(s); err != nil {
|
|
s.span.RecordError(err)
|
|
s.Errorf("done callback error: %s", err)
|
|
}
|
|
s.span.SetAttributes(config.AttributePeerTaskSuccess.Bool(true))
|
|
})
|
|
return nil
|
|
}
|
|
|
|
func (s *streamPeerTask) cleanUnfinished() {
|
|
// send last progress
|
|
s.once.Do(func() {
|
|
// send EOF piece result to scheduler
|
|
_ = s.peerPacketStream.Send(
|
|
scheduler.NewEndPieceResult(s.taskID, s.peerID, s.readyPieces.Settled()))
|
|
s.Debugf("end piece result sent")
|
|
close(s.done)
|
|
//close(s.successPieceCh)
|
|
if err := s.callback.Fail(s, s.failedCode, s.failedReason); err != nil {
|
|
s.span.RecordError(err)
|
|
s.Errorf("fail callback error: %s", err)
|
|
}
|
|
s.span.SetAttributes(config.AttributePeerTaskSuccess.Bool(false))
|
|
s.span.SetAttributes(config.AttributePeerTaskCode.Int(int(s.failedCode)))
|
|
s.span.SetAttributes(config.AttributePeerTaskMessage.String(s.failedReason))
|
|
})
|
|
}
|
|
|
|
func (s *streamPeerTask) SetContentLength(i int64) error {
|
|
s.contentLength = i
|
|
if !s.isCompleted() {
|
|
return nil
|
|
}
|
|
return s.finish()
|
|
}
|
|
|
|
func (s *streamPeerTask) writeTo(w io.Writer, pieceNum int32) (int64, error) {
|
|
pr, pc, err := s.pieceManager.ReadPiece(s.ctx, &storage.ReadPieceRequest{
|
|
PeerTaskMetaData: storage.PeerTaskMetaData{
|
|
PeerID: s.peerID,
|
|
TaskID: s.taskID,
|
|
},
|
|
PieceMetaData: storage.PieceMetaData{
|
|
Num: pieceNum,
|
|
},
|
|
})
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
n, err := io.Copy(w, pr)
|
|
if err != nil {
|
|
pc.Close()
|
|
return n, err
|
|
}
|
|
return n, pc.Close()
|
|
}
|