494 lines
16 KiB
Go
494 lines
16 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"
|
|
|
|
"github.com/go-http-utils/headers"
|
|
"github.com/pkg/errors"
|
|
"go.opentelemetry.io/otel/semconv"
|
|
"go.opentelemetry.io/otel/trace"
|
|
"go.uber.org/atomic"
|
|
"golang.org/x/time/rate"
|
|
|
|
"d7y.io/dragonfly/v2/client/config"
|
|
"d7y.io/dragonfly/v2/client/daemon/storage"
|
|
logger "d7y.io/dragonfly/v2/internal/dflog"
|
|
"d7y.io/dragonfly/v2/internal/idgen"
|
|
"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 starts the special peer task, return an io.Reader for stream io
|
|
// when all data transferred, reader return an io.EOF
|
|
// attribute stands some extra data, like HTTP response Header
|
|
Start(ctx context.Context) (rc io.ReadCloser, attribute map[string]string, err error)
|
|
}
|
|
|
|
type streamPeerTask struct {
|
|
peerTask
|
|
streamDone chan struct{}
|
|
successPieceCh chan int32
|
|
}
|
|
|
|
var _ StreamPeerTask = (*streamPeerTask)(nil)
|
|
|
|
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, pid: %s, url: %s, filter: %s, meta: %s, tag: %s",
|
|
request.PeerId, request.Url, request.UrlMeta.Filter, request.UrlMeta, request.UrlMeta.Tag)
|
|
// trace register
|
|
regCtx, cancel := context.WithTimeout(ctx, schedulerOption.ScheduleTimeout.Duration)
|
|
defer cancel()
|
|
regCtx, regSpan := tracer.Start(regCtx, config.SpanRegisterTask)
|
|
logger.Infof("step 1: peer %s start to register", request.PeerId)
|
|
result, err := schedulerClient.RegisterPeerTask(regCtx, request)
|
|
regSpan.RecordError(err)
|
|
regSpan.End()
|
|
|
|
var needBackSource bool
|
|
if err != nil {
|
|
if err == context.DeadlineExceeded {
|
|
logger.Errorf("scheduler did not response in %s", schedulerOption.ScheduleTimeout.Duration)
|
|
}
|
|
logger.Errorf("step 1: peer %s register failed: %s", request.PeerId, err)
|
|
if schedulerOption.DisableAutoBackSource {
|
|
logger.Errorf("register peer task failed: %s, peer id: %s, auto back source disabled", err, request.PeerId)
|
|
span.RecordError(err)
|
|
span.End()
|
|
return ctx, nil, nil, err
|
|
}
|
|
needBackSource = true
|
|
// can not detect source or scheduler error, create a new dummy scheduler client
|
|
schedulerClient = &dummySchedulerClient{}
|
|
result = &scheduler.RegisterResult{TaskId: idgen.TaskID(request.Url, request.UrlMeta)}
|
|
logger.Warnf("register peer task failed: %s, peer id: %s, try to back source", err, request.PeerId)
|
|
}
|
|
|
|
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 !needBackSource {
|
|
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)
|
|
logger.Infof("step 2: start report peer %s piece result", request.PeerId)
|
|
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))
|
|
}
|
|
pt := &streamPeerTask{
|
|
successPieceCh: make(chan int32),
|
|
streamDone: make(chan struct{}),
|
|
peerTask: peerTask{
|
|
ctx: ctx,
|
|
host: host,
|
|
needBackSource: needBackSource,
|
|
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,
|
|
readyPieces: NewBitmap(),
|
|
requestedPieces: NewBitmap(),
|
|
failedPieceCh: make(chan int32, config.DefaultPieceChanSize),
|
|
failedReason: failedReasonNotSet,
|
|
failedCode: base.Code_UnknownError,
|
|
contentLength: atomic.NewInt64(-1),
|
|
pieceParallelCount: atomic.NewInt32(0),
|
|
totalPiece: -1,
|
|
schedulerOption: schedulerOption,
|
|
schedulerClient: schedulerClient,
|
|
limiter: limiter,
|
|
completedLength: atomic.NewInt64(0),
|
|
usedTraffic: atomic.NewInt64(0),
|
|
SugaredLoggerOnWith: logger.With("peer", request.PeerId, "task", result.TaskId, "component", "streamPeerTask"),
|
|
},
|
|
}
|
|
// bind func that base peer task did not implement
|
|
pt.backSourceFunc = pt.backSource
|
|
pt.setContentLengthFunc = pt.SetContentLength
|
|
pt.setTotalPiecesFunc = pt.SetTotalPieces
|
|
pt.reportPieceResultFunc = pt.ReportPieceResult
|
|
return ctx, pt, nil, nil
|
|
}
|
|
|
|
func (s *streamPeerTask) ReportPieceResult(result *pieceTaskResult) error {
|
|
defer s.recoverFromPanic()
|
|
// retry failed piece
|
|
if !result.pieceResult.Success {
|
|
result.pieceResult.FinishedCount = s.readyPieces.Settled()
|
|
_ = s.peerPacketStream.Send(result.pieceResult)
|
|
if result.notRetry {
|
|
s.Warnf("piece %d download failed, no retry", result.piece.PieceNum)
|
|
return nil
|
|
}
|
|
select {
|
|
case <-s.done:
|
|
s.Infof("peer task done, stop to send failed piece")
|
|
case <-s.ctx.Done():
|
|
s.Debugf("context done due to %s, stop to send failed piece", s.ctx.Err())
|
|
case s.failedPieceCh <- result.pieceResult.PieceInfo.PieceNum:
|
|
s.Warnf("piece %d download failed, retry later", result.piece.PieceNum)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
s.lock.Lock()
|
|
if s.readyPieces.IsSet(result.pieceResult.PieceInfo.PieceNum) {
|
|
s.lock.Unlock()
|
|
s.Warnf("piece %d is already reported, skipped", result.pieceResult.PieceInfo.PieceNum)
|
|
return nil
|
|
}
|
|
// mark piece processed
|
|
s.readyPieces.Set(result.pieceResult.PieceInfo.PieceNum)
|
|
s.completedLength.Add(int64(result.piece.RangeSize))
|
|
s.lock.Unlock()
|
|
|
|
result.pieceResult.FinishedCount = s.readyPieces.Settled()
|
|
_ = s.peerPacketStream.Send(result.pieceResult)
|
|
s.successPieceCh <- result.piece.PieceNum
|
|
s.Debugf("success piece %d sent", result.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.ReadCloser, map[string]string, error) {
|
|
s.ctx, s.cancel = context.WithCancel(ctx)
|
|
if s.needBackSource {
|
|
go s.backSource()
|
|
} else {
|
|
s.pullPieces(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 != failedReasonNotSet {
|
|
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()
|
|
attr := map[string]string{}
|
|
attr[config.HeaderDragonflyTask] = s.taskID
|
|
attr[config.HeaderDragonflyPeer] = s.peerID
|
|
return nil, attr, err
|
|
case <-s.done:
|
|
var err error
|
|
if s.failedReason != failedReasonNotSet {
|
|
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()
|
|
attr := map[string]string{}
|
|
attr[config.HeaderDragonflyTask] = s.taskID
|
|
attr[config.HeaderDragonflyPeer] = s.peerID
|
|
return nil, attr, err
|
|
case first := <-s.successPieceCh:
|
|
firstPiece = first
|
|
}
|
|
|
|
attr := map[string]string{}
|
|
if s.contentLength.Load() != -1 {
|
|
attr[headers.ContentLength] = fmt.Sprintf("%d", s.contentLength.Load())
|
|
} else {
|
|
attr[headers.TransferEncoding] = "chunked"
|
|
}
|
|
attr[config.HeaderDragonflyTask] = s.taskID
|
|
attr[config.HeaderDragonflyPeer] = s.peerID
|
|
|
|
pr, pw := io.Pipe()
|
|
var readCloser io.ReadCloser = pr
|
|
go s.writeToPipe(firstPiece, pw)
|
|
|
|
return readCloser, attr, nil
|
|
}
|
|
|
|
func (s *streamPeerTask) finish() error {
|
|
// send last progress
|
|
s.once.Do(func() {
|
|
s.success = true
|
|
if err := s.callback.Update(s); err != nil {
|
|
s.span.RecordError(err)
|
|
s.Errorf("update callback error: %s", err)
|
|
}
|
|
// let stream return immediately
|
|
close(s.streamDone)
|
|
// send EOF piece result to scheduler
|
|
_ = s.peerPacketStream.Send(
|
|
scheduler.NewEndPieceResult(s.taskID, s.peerID, s.readyPieces.Settled()))
|
|
s.Debugf("end piece result sent, peer task finished")
|
|
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))
|
|
close(s.done)
|
|
})
|
|
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.Errorf("end piece result sent, peer task failed")
|
|
//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)
|
|
}
|
|
close(s.streamDone)
|
|
close(s.done)
|
|
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.Store(i)
|
|
if !s.isCompleted() {
|
|
return nil
|
|
}
|
|
return s.finish()
|
|
}
|
|
|
|
func (s *streamPeerTask) SetTotalPieces(i int32) {
|
|
s.totalPiece = i
|
|
}
|
|
|
|
func (s *streamPeerTask) writeOnePiece(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()
|
|
}
|
|
|
|
func (s *streamPeerTask) backSource() {
|
|
backSourceCtx, backSourceSpan := tracer.Start(s.ctx, config.SpanBackSource)
|
|
defer backSourceSpan.End()
|
|
s.contentLength.Store(-1)
|
|
_ = s.callback.Init(s)
|
|
reportPieceCtx, reportPieceSpan := tracer.Start(backSourceCtx, config.SpanReportPieceResult)
|
|
defer reportPieceSpan.End()
|
|
if peerPacketStream, err := s.schedulerClient.ReportPieceResult(reportPieceCtx, s.taskID, s.request); err != nil {
|
|
logger.Errorf("step 2: peer %s report piece failed: err", s.request.PeerId, err)
|
|
} else {
|
|
s.peerPacketStream = peerPacketStream
|
|
}
|
|
logger.Infof("step 2: start report peer %s back source piece result", s.request.PeerId)
|
|
err := s.pieceManager.DownloadSource(s.ctx, s, s.request)
|
|
if err != nil {
|
|
s.Errorf("download from source error: %s", err)
|
|
s.failedReason = err.Error()
|
|
s.cleanUnfinished()
|
|
backSourceSpan.SetAttributes(config.AttributePeerTaskSuccess.Bool(false))
|
|
backSourceSpan.RecordError(err)
|
|
return
|
|
}
|
|
s.Debugf("download from source ok")
|
|
backSourceSpan.SetAttributes(config.AttributePeerTaskSuccess.Bool(true))
|
|
_ = s.finish()
|
|
return
|
|
}
|
|
|
|
func (s *streamPeerTask) writeToPipe(firstPiece int32, pw *io.PipeWriter) {
|
|
defer func() {
|
|
s.cancel()
|
|
s.span.End()
|
|
}()
|
|
var (
|
|
desired int32
|
|
cur int32
|
|
wrote int64
|
|
err error
|
|
cache = make(map[int32]bool)
|
|
)
|
|
// update first piece to cache and check cur with desired
|
|
cache[firstPiece] = true
|
|
cur = firstPiece
|
|
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.writeOnePiece(pw, desired)
|
|
if err != nil {
|
|
span.RecordError(err)
|
|
span.End()
|
|
s.Errorf("write to pipe error: %s", err)
|
|
_ = pw.CloseWithError(err)
|
|
return
|
|
}
|
|
span.SetAttributes(config.AttributePieceSize.Int(int(wrote)))
|
|
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.streamDone:
|
|
for {
|
|
// all data wrote to local storage, and all data wrote to pipe write
|
|
if s.readyPieces.Settled() == desired {
|
|
if err = s.callback.ValidateDigest(s); err != nil {
|
|
s.span.RecordError(err)
|
|
s.Errorf("validate digest error: %s", err)
|
|
_ = pw.CloseWithError(err)
|
|
return
|
|
}
|
|
s.Debugf("all %d pieces wrote to pipe", desired)
|
|
pw.Close()
|
|
return
|
|
}
|
|
_, span := tracer.Start(s.ctx, config.SpanWriteBackPiece)
|
|
span.SetAttributes(config.AttributePiece.Int(int(desired)))
|
|
wrote, err = s.writeOnePiece(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.SetAttributes(config.AttributePieceSize.Int(int(wrote)))
|
|
span.End()
|
|
s.Debugf("wrote piece %d to pipe, size %d", desired, wrote)
|
|
desired++
|
|
}
|
|
case cur = <-s.successPieceCh:
|
|
continue
|
|
}
|
|
}
|
|
}
|