|
- // Copyright 2017 Google LLC
- //
- // 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 rpcreplay
-
- import (
- "io"
- "log"
- "net"
-
- "golang.org/x/net/context"
- "google.golang.org/grpc"
- "google.golang.org/grpc/codes"
- "google.golang.org/grpc/status"
-
- pb "cloud.google.com/go/rpcreplay/proto/intstore"
- )
-
- // intStoreServer is an in-memory implementation of IntStore.
- type intStoreServer struct {
- pb.IntStoreServer
-
- Addr string
- l net.Listener
- gsrv *grpc.Server
-
- items map[string]int32
- }
-
- func newIntStoreServer() *intStoreServer {
- l, err := net.Listen("tcp", "127.0.0.1:0")
- if err != nil {
- log.Fatal(err)
- }
- s := &intStoreServer{
- Addr: l.Addr().String(),
- l: l,
- gsrv: grpc.NewServer(),
- }
- pb.RegisterIntStoreServer(s.gsrv, s)
- go s.gsrv.Serve(s.l)
- return s
- }
-
- func (s *intStoreServer) stop() {
- s.gsrv.Stop()
- s.l.Close()
- }
-
- func (s *intStoreServer) Set(_ context.Context, item *pb.Item) (*pb.SetResponse, error) {
- old := s.setItem(item)
- return &pb.SetResponse{PrevValue: old}, nil
- }
-
- func (s *intStoreServer) setItem(item *pb.Item) int32 {
- if s.items == nil {
- s.items = map[string]int32{}
- }
- old := s.items[item.Name]
- s.items[item.Name] = item.Value
- return old
- }
-
- func (s *intStoreServer) Get(_ context.Context, req *pb.GetRequest) (*pb.Item, error) {
- val, ok := s.items[req.Name]
- if !ok {
- return nil, status.Errorf(codes.NotFound, "%q", req.Name)
- }
- return &pb.Item{Name: req.Name, Value: val}, nil
- }
-
- func (s *intStoreServer) ListItems(_ *pb.ListItemsRequest, ss pb.IntStore_ListItemsServer) error {
- for name, val := range s.items {
- if err := ss.Send(&pb.Item{Name: name, Value: val}); err != nil {
- return err
- }
- }
- return nil
- }
-
- func (s *intStoreServer) SetStream(ss pb.IntStore_SetStreamServer) error {
- n := 0
- for {
- item, err := ss.Recv()
- if err == io.EOF {
- break
- }
- if err != nil {
- return err
- }
- s.setItem(item)
- n++
- }
- return ss.SendAndClose(&pb.Summary{Count: int32(n)})
- }
-
- func (s *intStoreServer) StreamChat(ss pb.IntStore_StreamChatServer) error {
- for {
- item, err := ss.Recv()
- if err == io.EOF {
- break
- }
- if err != nil {
- return err
- }
- if err := ss.Send(item); err != nil {
- return err
- }
- }
- return nil
- }
|