go.uber.org/yarpc@v1.72.1/x/yarpctest/types/handler_unary.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package types
    22  
    23  import (
    24  	"context"
    25  	"fmt"
    26  	"testing"
    27  
    28  	"go.uber.org/atomic"
    29  	"go.uber.org/multierr"
    30  	"go.uber.org/yarpc/api/middleware"
    31  	"go.uber.org/yarpc/api/transport"
    32  	"go.uber.org/yarpc/internal/inboundmiddleware"
    33  	"go.uber.org/yarpc/x/yarpctest/api"
    34  )
    35  
    36  // UnaryHandler is a struct that implements the ProcOptions and HandlerOption
    37  // interfaces (so it can be used directly as a procedure, or as a single use
    38  // handler (depending on the use case)).
    39  type UnaryHandler struct {
    40  	Handler    api.UnaryHandler
    41  	Middleware []api.UnaryInboundMiddleware
    42  }
    43  
    44  // Start implements Lifecycle.
    45  func (h *UnaryHandler) Start(t testing.TB) error {
    46  	var err error
    47  	err = multierr.Append(err, h.Handler.Start(t))
    48  	for _, mw := range h.Middleware {
    49  		err = multierr.Append(err, mw.Start(t))
    50  	}
    51  	return err
    52  }
    53  
    54  // Stop implements Lifecycle.
    55  func (h *UnaryHandler) Stop(t testing.TB) error {
    56  	var err error
    57  	err = multierr.Append(err, h.Handler.Stop(t))
    58  	for _, mw := range h.Middleware {
    59  		err = multierr.Append(err, mw.Stop(t))
    60  	}
    61  	return err
    62  }
    63  
    64  // ApplyProc implements ProcOption.
    65  func (h *UnaryHandler) ApplyProc(opts *api.ProcOpts) {
    66  	opts.HandlerSpec = transport.NewUnaryHandlerSpec(h)
    67  }
    68  
    69  // ApplyHandler implements HandlerOption.
    70  func (h *UnaryHandler) ApplyHandler(opts *api.HandlerOpts) {
    71  	opts.Handlers = append(opts.Handlers, h)
    72  }
    73  
    74  // Handle implements transport.UnaryHandler.
    75  func (h *UnaryHandler) Handle(ctx context.Context, req *transport.Request, resw transport.ResponseWriter) error {
    76  	mws := make([]middleware.UnaryInbound, 0, len(h.Middleware))
    77  	for _, mw := range h.Middleware {
    78  		mws = append(mws, mw)
    79  	}
    80  	handler := middleware.ApplyUnaryInbound(h.Handler, inboundmiddleware.UnaryChain(mws...))
    81  	return handler.Handle(ctx, req, resw)
    82  }
    83  
    84  // OrderedHandler implements the transport.UnaryHandler and ProcOption
    85  // interfaces.
    86  type OrderedHandler struct {
    87  	attempt  atomic.Int32
    88  	Handlers []api.UnaryHandler
    89  }
    90  
    91  // Start implements Lifecycle.
    92  func (h *OrderedHandler) Start(t testing.TB) error {
    93  	var err error
    94  	for _, handler := range h.Handlers {
    95  		err = multierr.Append(err, handler.Start(t))
    96  	}
    97  	return err
    98  }
    99  
   100  // Stop implements Lifecycle.
   101  func (h *OrderedHandler) Stop(t testing.TB) error {
   102  	var err error
   103  	for _, handler := range h.Handlers {
   104  		err = multierr.Append(err, handler.Stop(t))
   105  	}
   106  	return err
   107  }
   108  
   109  // ApplyProc implements ProcOption.
   110  func (h *OrderedHandler) ApplyProc(opts *api.ProcOpts) {
   111  	opts.HandlerSpec = transport.NewUnaryHandlerSpec(h)
   112  }
   113  
   114  // Handle implements transport.UnaryHandler#Handle.
   115  func (h *OrderedHandler) Handle(ctx context.Context, req *transport.Request, resw transport.ResponseWriter) error {
   116  	if len(h.Handlers) <= 0 {
   117  		return fmt.Errorf("no handlers for the request")
   118  	}
   119  	n := h.attempt.Inc()
   120  	if int(n) > len(h.Handlers) {
   121  		return fmt.Errorf("too many requests, expected %d, got %d", len(h.Handlers), n)
   122  	}
   123  	return h.Handlers[n-1].Handle(ctx, req, resw)
   124  }