main.go

Documentation: github.com/containerd/ttrpc/example/cmd

     1  /*
     2     Copyright The containerd Authors.
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package main
    18  
    19  import (
    20  	context "context"
    21  	"encoding/json"
    22  	"errors"
    23  	"log"
    24  	"net"
    25  	"os"
    26  
    27  	ttrpc "github.com/containerd/ttrpc"
    28  	"github.com/containerd/ttrpc/example"
    29  	"google.golang.org/protobuf/types/known/emptypb"
    30  )
    31  
    32  const socket = "example-ttrpc-server"
    33  
    34  func main() {
    35  	if err := handle(); err != nil {
    36  		log.Fatal(err)
    37  	}
    38  }
    39  
    40  func handle() error {
    41  	command := os.Args[1]
    42  	switch command {
    43  	case "server":
    44  		return server()
    45  	case "client":
    46  		return client()
    47  	default:
    48  		return errors.New("invalid command")
    49  	}
    50  }
    51  
    52  func serverIntercept(ctx context.Context, um ttrpc.Unmarshaler, i *ttrpc.UnaryServerInfo, m ttrpc.Method) (interface{}, error) {
    53  	log.Println("server interceptor")
    54  	dumpMetadata(ctx)
    55  	return m(ctx, um)
    56  }
    57  
    58  func clientIntercept(ctx context.Context, req *ttrpc.Request, resp *ttrpc.Response, i *ttrpc.UnaryClientInfo, invoker ttrpc.Invoker) error {
    59  	log.Println("client interceptor")
    60  	dumpMetadata(ctx)
    61  	return invoker(ctx, req, resp)
    62  }
    63  
    64  func dumpMetadata(ctx context.Context) {
    65  	md, ok := ttrpc.GetMetadata(ctx)
    66  	if !ok {
    67  		panic("no metadata")
    68  	}
    69  	if err := json.NewEncoder(os.Stdout).Encode(md); err != nil {
    70  		panic(err)
    71  	}
    72  }
    73  
    74  func server() error {
    75  	s, err := ttrpc.NewServer(
    76  		ttrpc.WithServerHandshaker(defaultHandshaker()),
    77  		ttrpc.WithUnaryServerInterceptor(serverIntercept),
    78  	)
    79  	if err != nil {
    80  		return err
    81  	}
    82  	defer s.Close()
    83  	example.RegisterExampleService(s, &exampleServer{})
    84  
    85  	l, err := net.Listen("unix", socket)
    86  	if err != nil {
    87  		return err
    88  	}
    89  	defer func() {
    90  		l.Close()
    91  		os.Remove(socket)
    92  	}()
    93  	return s.Serve(context.Background(), l)
    94  }
    95  
    96  func client() error {
    97  	conn, err := net.Dial("unix", socket)
    98  	if err != nil {
    99  		return err
   100  	}
   101  	defer conn.Close()
   102  
   103  	tc := ttrpc.NewClient(conn, ttrpc.WithUnaryClientInterceptor(clientIntercept))
   104  	client := example.NewExampleClient(tc)
   105  
   106  	r := &example.Method1Request{
   107  		Foo: os.Args[2],
   108  		Bar: os.Args[3],
   109  	}
   110  
   111  	ctx := context.Background()
   112  	md := ttrpc.MD{}
   113  	md.Set("name", "koye")
   114  	ctx = ttrpc.WithMetadata(ctx, md)
   115  
   116  	resp, err := client.Method1(ctx, r)
   117  	if err != nil {
   118  		return err
   119  	}
   120  	return json.NewEncoder(os.Stdout).Encode(resp)
   121  }
   122  
   123  type exampleServer struct {
   124  }
   125  
   126  func (s *exampleServer) Method1(ctx context.Context, r *example.Method1Request) (*example.Method1Response, error) {
   127  	return &example.Method1Response{
   128  		Foo: r.Foo,
   129  		Bar: r.Bar,
   130  	}, nil
   131  }
   132  
   133  func (s *exampleServer) Method2(ctx context.Context, r *example.Method1Request) (*emptypb.Empty, error) {
   134  	return &emptypb.Empty{}, nil
   135  }
   136
View as plain text