feat: implement PreheatFile and StatFile for scheduler (#4380) (#4381)

Co-authored-by: caoyifan.cyf <[email protected]>

Author: EvanCley

go.mod

@@ -4,7 +4,7 @@ go 1.23.8
 
 require (
 	cloud.google.com/go/storage v1.50.0
-	d7y.io/api/v2 v2.1.79
+	d7y.io/api/v2 v2.1.80
 	github.com/MysteriousPotato/go-lockable v1.0.0
 	github.com/Showmax/go-fqdn v1.0.0
 	github.com/VividCortex/mysqlerr v1.0.0

go.sum

@@ -63,8 +63,8 @@ cloud.google.com/go/storage v1.50.0 h1:3TbVkzTooBvnZsk7WaAQfOsNrdoM8QHusXA1cpk6Q
 cloud.google.com/go/storage v1.50.0/go.mod h1:l7XeiD//vx5lfqE3RavfmU9yvk5Pp0Zhcv482poyafY=
 cloud.google.com/go/trace v1.11.6 h1:2O2zjPzqPYAHrn3OKl029qlqG6W8ZdYaOWRyr8NgMT4=
 cloud.google.com/go/trace v1.11.6/go.mod h1:GA855OeDEBiBMzcckLPE2kDunIpC72N+Pq8WFieFjnI=
-d7y.io/api/v2 v2.1.79 h1:LUPpjiuhHZgyezHBH7lG42pkqjcdbiTAHtEJ2Z7tUJg=
-d7y.io/api/v2 v2.1.79/go.mod h1:t6k27g8dFyH6sp3y2J1qHyk2YmoySd88qSbsEaqJ2Sw=
+d7y.io/api/v2 v2.1.80 h1:dH4w0dUXLpZ1sdp6QuWUVXkjQaUSX/pAOFFunkrcsUs=
+d7y.io/api/v2 v2.1.80/go.mod h1:t6k27g8dFyH6sp3y2J1qHyk2YmoySd88qSbsEaqJ2Sw=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
 filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
 filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=

internal/dflog/logger.go

@@ -256,6 +256,24 @@ func WithStatImageAndTaskID(url, taskID string) *SugaredLoggerOnWith {
 	}
 }
 
+func WithPreheatFile(url string) *SugaredLoggerOnWith {
+	return &SugaredLoggerOnWith{
+		withArgs: []any{"url", url},
+	}
+}
+
+func WithStatFile(url string) *SugaredLoggerOnWith {
+	return &SugaredLoggerOnWith{
+		withArgs: []any{"url", url},
+	}
+}
+
+func WithStatFileAndTaskID(url, taskID string) *SugaredLoggerOnWith {
+	return &SugaredLoggerOnWith{
+		withArgs: []any{"url", url, "taskID", taskID},
+	}
+}
+
 func (log *SugaredLoggerOnWith) With(args ...any) *SugaredLoggerOnWith {
 	args = append(args, log.withArgs...)
 	return &SugaredLoggerOnWith{

internal/job/types.go

@@ -18,6 +18,11 @@ package job
 
 import (
 	"time"
+
+	"google.golang.org/protobuf/types/known/durationpb"
+
+	v2 "d7y.io/api/v2/pkg/apis/common/v2"
+	dfdaemonv2 "d7y.io/api/v2/pkg/apis/dfdaemon/v2"
 )
 
 // PreheatRequest defines the request parameters for preheating.
@@ -83,6 +88,24 @@ type GetTaskResponse struct {
 	SchedulerClusterID uint    `json:"scheduler_cluster_id"`
 }
 
+// ListTaskEntriesRequest defines the request parameters for listing task entries.
+type ListTaskEntriesRequest struct {
+	TaskID           string               `json:"task_id" validate:"required"`
+	Url              string               `json:"url" validate:"omitempty"`
+	Timeout          *durationpb.Duration `json:"timeout" validate:"omitempty"`
+	Header           map[string]string    `json:"header" validate:"omitempty"`
+	CertificateChain [][]byte             `json:"certificate_chain" validate:"omitempty"`
+	ObjectStorage    *v2.ObjectStorage    `json:"object_storage" validate:"omitempty"`
+	Hdfs             *v2.HDFS             `json:"hdfs" validate:"omitempty"`
+}
+
+// ListTaskEntriesResponse defines the response parameters for listing task entries.
+type ListTaskEntriesResponse struct {
+	Entries     []*dfdaemonv2.Entry `json:"entries"`
+	Recursive   bool                `json:"recursive"`
+	SchedulerID uint                `json:"scheduler_id"`
+}
+
 // Peer represents the peer information.
 type Peer struct {
 	ID        string    `json:"id"`

pkg/rpc/dfdaemon/client/client_v2.go

@@ -89,6 +89,9 @@ type V2 interface {
 	// DeleteTask deletes task from p2p network.
 	DeleteTask(context.Context, *dfdaemonv2.DeleteTaskRequest, ...grpc.CallOption) error
 
+	// ListTaskEntries lists task entries.
+	ListTaskEntries(context.Context, *dfdaemonv2.ListTaskEntriesRequest, ...grpc.CallOption) (*dfdaemonv2.ListTaskEntriesResponse, error)
+
 	// DownloadPersistentCacheTask downloads persistent cache task from p2p network.
 	DownloadPersistentCacheTask(context.Context, *dfdaemonv2.DownloadPersistentCacheTaskRequest, ...grpc.CallOption) (dfdaemonv2.DfdaemonUpload_DownloadPersistentCacheTaskClient, error)
 
@@ -192,3 +195,11 @@ func (v *v2) DeletePersistentCacheTask(ctx context.Context, req *dfdaemonv2.Dele
 	_, err := v.DfdaemonUploadClient.DeletePersistentCacheTask(ctx, req, opts...)
 	return err
 }
+
+// ListTaskEntries lists task entries from p2p network.
+func (v *v2) ListTaskEntries(ctx context.Context, req *dfdaemonv2.ListTaskEntriesRequest, opts ...grpc.CallOption) (*dfdaemonv2.ListTaskEntriesResponse, error) {
+	ctx, cancel := context.WithTimeout(ctx, contextTimeout)
+	defer cancel()
+
+	return v.DfdaemonUploadClient.ListTaskEntries(ctx, req, opts...)
+}

scheduler/job/job.go

@@ -59,6 +59,9 @@ type Job interface {
 	// GetTask retrieves task information from all hosts in the cluster.
 	GetTask(context.Context, *internaljob.GetTaskRequest, *logger.SugaredLoggerOnWith) (*internaljob.GetTaskResponse, error)
 
+	// ListTaskEntries lists all task entries.
+	ListTaskEntries(context.Context, *internaljob.ListTaskEntriesRequest, *logger.SugaredLoggerOnWith) (*internaljob.ListTaskEntriesResponse, error)
+
 	// PreheatSinglePeer preheats job by single seed peer, scheduler will trigger seed peer to download task.
 	PreheatSingleSeedPeer(context.Context, *internaljob.PreheatRequest, *logger.SugaredLoggerOnWith) (*internaljob.PreheatResponse, error)
 
@@ -1034,3 +1037,47 @@ func (j *job) deleteTask(ctx context.Context, data string) (string, error) {
 		SchedulerClusterID: j.config.Manager.SchedulerClusterID,
 	})
 }
+
+func (j *job) ListTaskEntries(ctx context.Context, req *internaljob.ListTaskEntriesRequest, log *logger.SugaredLoggerOnWith) (*internaljob.ListTaskEntriesResponse, error) {
+	advertiseIP := j.config.Server.AdvertiseIP.String()
+
+	selected, err := j.resource.SeedPeer().Select(ctx, req.TaskID)
+	if err != nil {
+		return nil, err
+	}
+
+	addr := fmt.Sprintf("%s:%d", selected.IP, selected.Port)
+	log.Infof("selected seed peer %s for task %s", addr, req.TaskID)
+
+	dfdaemonClient, err := dfdaemonclient.GetV2ByAddr(ctx, addr, j.dialOptions...)
+	if err != nil {
+		log.Errorf("[list-task-entries] get dfdaemon client failed: %s", err)
+		return nil, err
+	}
+
+	res, err := dfdaemonClient.ListTaskEntries(ctx, &dfdaemonv2.ListTaskEntriesRequest{
+		TaskId:           req.TaskID,
+		Url:              req.Url,
+		RequestHeader:    req.Header,
+		Timeout:          req.Timeout,
+		CertificateChain: req.CertificateChain,
+		ObjectStorage:    req.ObjectStorage,
+		Hdfs:             req.Hdfs,
+		RemoteIp:         &advertiseIP,
+	})
+	if err != nil {
+		log.Errorf("[list-task-entries]list task entries failed: %s", err)
+		return nil, err
+	}
+
+	var recursive bool
+	if len(res.Entries) > 1 {
+		recursive = true
+	}
+
+	return &internaljob.ListTaskEntriesResponse{
+		Recursive:   recursive,
+		Entries:     res.Entries,
+		SchedulerID: j.config.Manager.SchedulerClusterID,
+	}, nil
+}

scheduler/job/mocks/job_mock.go

@@ -485,6 +485,34 @@ var (
 		Help:      "Counter of the number of failed of the stating image.",
 	})
 
+	PreheatFileCount = promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: types.MetricsNamespace,
+		Subsystem: types.SchedulerMetricsName,
+		Name:      "preheat_file_total",
+		Help:      "Counter of the total preheat file.",
+	})
+
+	PreheatFileFailureCount = promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: types.MetricsNamespace,
+		Subsystem: types.SchedulerMetricsName,
+		Name:      "preheat_file_failure_total",
+		Help:      "Counter of the total preheat file failure.",
+	})
+
+	StatFileCount = promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: types.MetricsNamespace,
+		Subsystem: types.SchedulerMetricsName,
+		Name:      "stat_file_total",
+		Help:      "Counter of the number of the stating file.",
+	})
+
+	StatFileFailureCount = promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: types.MetricsNamespace,
+		Subsystem: types.SchedulerMetricsName,
+		Name:      "stat_file_failure_total",
+		Help:      "Counter of the number of failed of the stating file.",
+	})
+
 	VersionGauge = promauto.NewGaugeVec(prometheus.GaugeOpts{
 		Namespace: types.MetricsNamespace,
 		Subsystem: types.SchedulerMetricsName,

scheduler/metrics/metrics.go

@@ -342,12 +342,42 @@ func (s *schedulerServerV2) StatImage(ctx context.Context, req *schedulerv2.Stat
 	return resp, nil
 }
 
-// TODO(EvanCley): Implement the following methods.
+// PreheatFile synchronously triggers an asynchronous preheat task for a file.
+//
+// This is a blocking call. The RPC will not return until the server has completed the
+// initial synchronous work: preparing the file URL.
+//
+// After this call successfully returns, a scheduler on the server begins the actual
+// preheating process, instructing peers to download the file in the background.
+//
+// A successful response (google.protobuf.Empty) confirms that the preparation is complete
+// and the asynchronous download task has been scheduled.
 func (s *schedulerServerV2) PreheatFile(ctx context.Context, req *schedulerv2.PreheatFileRequest) (*emptypb.Empty, error) {
+	// Collect PreheatFileCount metrics.
+	metrics.PreheatFileCount.Inc()
+	if err := s.service.PreheatFile(ctx, req); err != nil {
+		// Collect PreheatFileFailureCount metrics.
+		metrics.PreheatFileFailureCount.Inc()
+		return nil, err
+	}
+
 	return new(emptypb.Empty), nil
 }
 
-// TODO(EvanCley): Implement the following methods.
+// StatFile provides detailed status for files distribution in peers.
+//
+// This is a blocking call that first queries the file/dir entries and then queries
+// all peers to collect the file's download state across the network.
+// The response includes the file status on each peer.
 func (s *schedulerServerV2) StatFile(ctx context.Context, req *schedulerv2.StatFileRequest) (*schedulerv2.StatFileResponse, error) {
-	return nil, nil
+	// Collect StatFileCount metrics.
+	metrics.StatFileCount.Inc()
+	resp, err := s.service.StatFile(ctx, req)
+	if err != nil {
+		// Collect StatFileFailureCount metrics.
+		metrics.StatFileFailureCount.Inc()
+		return nil, err
+	}
+
+	return resp, nil
 }