-
Notifications
You must be signed in to change notification settings - Fork 281
/
rpcserver.go
111 lines (92 loc) · 3.59 KB
/
rpcserver.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
/*
* 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 rpcserver
import (
"context"
"google.golang.org/grpc"
"google.golang.org/grpc/health"
healthpb "google.golang.org/grpc/health/grpc_health_v1"
empty "google.golang.org/protobuf/types/known/emptypb"
"d7y.io/dragonfly/v2/pkg/rpc"
"d7y.io/dragonfly/v2/pkg/rpc/scheduler"
"d7y.io/dragonfly/v2/scheduler/metrics"
"d7y.io/dragonfly/v2/scheduler/resource"
"d7y.io/dragonfly/v2/scheduler/service"
)
// Server is grpc server.
type Server struct {
// Service interface.
service *service.Service
// GRPC UnimplementedSchedulerServer interface.
scheduler.UnimplementedSchedulerServer
}
// New returns a new transparent scheduler server from the given options.
func New(service *service.Service, opts ...grpc.ServerOption) *grpc.Server {
svr := &Server{service: service}
grpcServer := grpc.NewServer(append(rpc.DefaultServerOptions(), opts...)...)
// Register servers on grpc server.
scheduler.RegisterSchedulerServer(grpcServer, svr)
healthpb.RegisterHealthServer(grpcServer, health.NewServer())
return grpcServer
}
// RegisterPeerTask registers peer and triggers seed peer download task.
func (s *Server) RegisterPeerTask(ctx context.Context, req *scheduler.PeerTaskRequest) (*scheduler.RegisterResult, error) {
bizTag := resource.DefaultBizTag
if req.UrlMeta.Tag != "" {
bizTag = req.UrlMeta.Tag
}
metrics.RegisterPeerTaskCount.WithLabelValues(bizTag).Inc()
resp, err := s.service.RegisterPeerTask(ctx, req)
if err != nil {
metrics.RegisterPeerTaskFailureCount.WithLabelValues(bizTag).Inc()
} else {
metrics.PeerTaskCounter.WithLabelValues(bizTag, resp.SizeScope.String()).Inc()
}
return resp, err
}
// ReportPieceResult handles the piece information reported by dfdaemon.
func (s *Server) ReportPieceResult(stream scheduler.Scheduler_ReportPieceResultServer) error {
metrics.ConcurrentScheduleGauge.Inc()
defer metrics.ConcurrentScheduleGauge.Dec()
return s.service.ReportPieceResult(stream)
}
// ReportPeerResult handles peer result reported by dfdaemon.
func (s *Server) ReportPeerResult(ctx context.Context, req *scheduler.PeerResult) (*empty.Empty, error) {
return new(empty.Empty), s.service.ReportPeerResult(ctx, req)
}
// StatTask checks if the given task exists.
func (s *Server) StatTask(ctx context.Context, req *scheduler.StatTaskRequest) (*scheduler.Task, error) {
metrics.StatTaskCount.Inc()
task, err := s.service.StatTask(ctx, req)
if err != nil {
metrics.StatTaskFailureCount.Inc()
return nil, err
}
return task, nil
}
// AnnounceTask informs scheduler a peer has completed task.
func (s *Server) AnnounceTask(ctx context.Context, req *scheduler.AnnounceTaskRequest) (*empty.Empty, error) {
metrics.AnnounceCount.Inc()
if err := s.service.AnnounceTask(ctx, req); err != nil {
metrics.AnnounceFailureCount.Inc()
return new(empty.Empty), err
}
return new(empty.Empty), nil
}
// LeaveTask makes the peer unschedulable.
func (s *Server) LeaveTask(ctx context.Context, req *scheduler.PeerTarget) (*empty.Empty, error) {
return new(empty.Empty), s.service.LeaveTask(ctx, req)
}