github.com/v2fly/v2ray-core/v4@v4.45.2/app/router/command/command.go (about)

     1  //go:build !confonly
     2  // +build !confonly
     3  
     4  package command
     5  
     6  //go:generate go run github.com/v2fly/v2ray-core/v4/common/errors/errorgen
     7  
     8  import (
     9  	"context"
    10  	"time"
    11  
    12  	"google.golang.org/grpc"
    13  
    14  	core "github.com/v2fly/v2ray-core/v4"
    15  	"github.com/v2fly/v2ray-core/v4/common"
    16  	"github.com/v2fly/v2ray-core/v4/features/routing"
    17  	"github.com/v2fly/v2ray-core/v4/features/stats"
    18  )
    19  
    20  // routingServer is an implementation of RoutingService.
    21  type routingServer struct {
    22  	router       routing.Router
    23  	routingStats stats.Channel
    24  }
    25  
    26  // NewRoutingServer creates a statistics service with statistics manager.
    27  func NewRoutingServer(router routing.Router, routingStats stats.Channel) RoutingServiceServer {
    28  	return &routingServer{
    29  		router:       router,
    30  		routingStats: routingStats,
    31  	}
    32  }
    33  
    34  func (s *routingServer) TestRoute(ctx context.Context, request *TestRouteRequest) (*RoutingContext, error) {
    35  	if request.RoutingContext == nil {
    36  		return nil, newError("Invalid routing request.")
    37  	}
    38  	route, err := s.router.PickRoute(AsRoutingContext(request.RoutingContext))
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	if request.PublishResult && s.routingStats != nil {
    43  		ctx, _ := context.WithTimeout(context.Background(), 4*time.Second) // nolint: govet
    44  		s.routingStats.Publish(ctx, route)
    45  	}
    46  	return AsProtobufMessage(request.FieldSelectors)(route), nil
    47  }
    48  
    49  func (s *routingServer) SubscribeRoutingStats(request *SubscribeRoutingStatsRequest, stream RoutingService_SubscribeRoutingStatsServer) error {
    50  	if s.routingStats == nil {
    51  		return newError("Routing statistics not enabled.")
    52  	}
    53  	genMessage := AsProtobufMessage(request.FieldSelectors)
    54  	subscriber, err := stats.SubscribeRunnableChannel(s.routingStats)
    55  	if err != nil {
    56  		return err
    57  	}
    58  	defer stats.UnsubscribeClosableChannel(s.routingStats, subscriber)
    59  	for {
    60  		select {
    61  		case value, ok := <-subscriber:
    62  			if !ok {
    63  				return newError("Upstream closed the subscriber channel.")
    64  			}
    65  			route, ok := value.(routing.Route)
    66  			if !ok {
    67  				return newError("Upstream sent malformed statistics.")
    68  			}
    69  			err := stream.Send(genMessage(route))
    70  			if err != nil {
    71  				return err
    72  			}
    73  		case <-stream.Context().Done():
    74  			return stream.Context().Err()
    75  		}
    76  	}
    77  }
    78  
    79  func (s *routingServer) mustEmbedUnimplementedRoutingServiceServer() {}
    80  
    81  type service struct {
    82  	v *core.Instance
    83  }
    84  
    85  func (s *service) Register(server *grpc.Server) {
    86  	common.Must(s.v.RequireFeatures(func(router routing.Router, stats stats.Manager) {
    87  		RegisterRoutingServiceServer(server, NewRoutingServer(router, nil))
    88  	}))
    89  }
    90  
    91  func init() {
    92  	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, cfg interface{}) (interface{}, error) {
    93  		s := core.MustFromContext(ctx)
    94  		return &service{v: s}, nil
    95  	}))
    96  }