github.com/brownsys/tracing-framework-go@v0.0.0-20161210174012-0542a62412fe/xtrace/grpc/example/main.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"os"
     7  
     8  	"github.com/brownsys/tracing-framework-go/xtrace/client"
     9  	xtgrpc "github.com/brownsys/tracing-framework-go/xtrace/grpc"
    10  	"github.com/spf13/pflag"
    11  	"golang.org/x/net/context"
    12  	"google.golang.org/grpc"
    13  )
    14  
    15  var (
    16  	xtraceServerFlag string
    17  	nextServerFlag   string
    18  	listenAddrFlag   string
    19  	addFlag          int64
    20  	subFlag          int64
    21  	mulFlag          int64
    22  	divFlag          int64
    23  	inputFlag        int64
    24  )
    25  
    26  func init() {
    27  	pflag.StringVar(&xtraceServerFlag, "xtrace-server", "", "The X-Trace server to send logs to.")
    28  	pflag.StringVar(&nextServerFlag, "next-server", "", "The next server to invoke RPCs on, if any.")
    29  	pflag.StringVar(&listenAddrFlag, "listen-addr", "", "The address to listen for RPC calls on.")
    30  	pflag.Int64Var(&addFlag, "add", 0, "Add this value to the argument.")
    31  	pflag.Int64Var(&subFlag, "sub", 0, "Subtract this value from the argument.")
    32  	pflag.Int64Var(&mulFlag, "mul", 0, "Multiply the argument by this value.")
    33  	pflag.Int64Var(&divFlag, "div", 0, "Divide the argument by this value.")
    34  	pflag.Int64Var(&inputFlag, "in", 0, "Instead of hosting an RPC server, call --next-server with this argument.")
    35  }
    36  
    37  type server struct {
    38  	op       func(int64) int64
    39  	opname   string
    40  	next     OperatorClient
    41  	nextname string
    42  }
    43  
    44  func (s *server) PerformOp(ctx context.Context, in *Int) (*Int, error) {
    45  	err := xtgrpc.ExtractIDs(ctx)
    46  	if err != nil {
    47  		client.Log(fmt.Sprintf("could not extract Task ID in RPC handler: %v", err))
    48  	}
    49  
    50  	out := &Int{s.op(in.Num)}
    51  	client.Log(fmt.Sprintf("applying operation %v to argument %v: %v", s.opname, in.Num, out.Num))
    52  	if s.next != nil {
    53  		client.Log(fmt.Sprintf("passing to %v", s.nextname))
    54  		out, err = s.next.PerformOp(ctx, out)
    55  		if err != nil {
    56  			client.Log(fmt.Sprintf("error invoking RPC: %v", err))
    57  			return nil, err
    58  		}
    59  		return out, nil
    60  	}
    61  	client.Log(fmt.Sprintf("returning value %v", out.Num))
    62  	return out, nil
    63  }
    64  
    65  func main() {
    66  	pflag.Parse()
    67  
    68  	usage := func() { pflag.PrintDefaults(); os.Exit(1) }
    69  
    70  	// number of operations flags set
    71  	numSet := 0
    72  	for _, f := range []string{"add", "sub", "mul", "div"} {
    73  		if pflag.Lookup(f).Changed {
    74  			numSet++
    75  		}
    76  	}
    77  
    78  	switch {
    79  	case !pflag.Lookup("xtrace-server").Changed:
    80  		fmt.Fprintln(os.Stderr, "Must provide --xtrace-server")
    81  		usage()
    82  	case pflag.Lookup("in").Changed && (numSet > 0 || pflag.Lookup("listen-addr").Changed):
    83  		fmt.Fprintln(os.Stderr, "Cannot provide --in and any of --listen-addr, --add, --sub, --mul, or --div")
    84  		usage()
    85  	case !pflag.Lookup("in").Changed && (numSet == 0 || numSet > 1 || !pflag.Lookup("listen-addr").Changed):
    86  		fmt.Fprintln(os.Stderr, "Must provide --in or exactly one of --add, --sub, --mul, --div, and also --listen-addr")
    87  		usage()
    88  	case pflag.Lookup("in").Changed && !pflag.Lookup("next-server").Changed:
    89  		fmt.Fprintln(os.Stderr, "Cannot provide --in without --next-server")
    90  		usage()
    91  	}
    92  
    93  	if err := client.Connect(xtraceServerFlag); err != nil {
    94  		fmt.Fprintf(os.Stderr, "could not connect to X-Trace server: %v\n", err)
    95  		os.Exit(2)
    96  	}
    97  
    98  	var next OperatorClient
    99  	if pflag.Lookup("next-server").Changed {
   100  		var err error
   101  		cc, err := grpc.Dial(nextServerFlag, grpc.WithInsecure())
   102  		if err != nil {
   103  			fmt.Fprintf(os.Stderr, "could not connect to next gRPC server: %v\n", err)
   104  			os.Exit(2)
   105  		}
   106  		next = NewOperatorClient(cc)
   107  	}
   108  
   109  	if pflag.Lookup("in").Changed {
   110  		client.Log(fmt.Sprintf("passing --in (%v) to %v", inputFlag, nextServerFlag))
   111  		out, err := next.PerformOp(context.Background(), &Int{inputFlag})
   112  		if err != nil {
   113  			client.Log(fmt.Sprintf("error invoking RPC: %v", err))
   114  			os.Exit(2)
   115  		}
   116  		fmt.Println("Got value", out.Num)
   117  		return
   118  	}
   119  
   120  	var op func(int64) int64
   121  	var opname string
   122  	switch {
   123  	case pflag.Lookup("add").Changed:
   124  		op = func(in int64) int64 { return in + addFlag }
   125  		opname = fmt.Sprintf("+%v", addFlag)
   126  	case pflag.Lookup("sub").Changed:
   127  		op = func(in int64) int64 { return in - subFlag }
   128  		opname = fmt.Sprintf("-%v", subFlag)
   129  	case pflag.Lookup("mul").Changed:
   130  		op = func(in int64) int64 { return in * mulFlag }
   131  		opname = fmt.Sprintf("*%v", mulFlag)
   132  	case pflag.Lookup("div").Changed:
   133  		op = func(in int64) int64 { return in / divFlag }
   134  		opname = fmt.Sprintf("/%v", divFlag)
   135  	}
   136  
   137  	srv := server{
   138  		op:       op,
   139  		opname:   opname,
   140  		next:     next,
   141  		nextname: nextServerFlag,
   142  	}
   143  
   144  	lis, err := net.Listen("tcp", listenAddrFlag)
   145  	if err != nil {
   146  		fmt.Fprintln(os.Stderr, "Failed to listen:", err)
   147  		os.Exit(2)
   148  	}
   149  	s := grpc.NewServer()
   150  	RegisterOperatorServer(s, &srv)
   151  	s.Serve(lis)
   152  }