github.com/ferranbt/nomad@v0.9.3-0.20190607002617-85c449b7667c/nomad/structs/streaming_rpc.go (about)

     1  package structs
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"sync"
     7  )
     8  
     9  // StreamingRpcHeader is the first struct serialized after entering the
    10  // streaming RPC mode. The header is used to dispatch to the correct method.
    11  type StreamingRpcHeader struct {
    12  	// Method is the name of the method to invoke.
    13  	Method string
    14  }
    15  
    16  // StreamingRpcAck is used to acknowledge receiving the StreamingRpcHeader and
    17  // routing to the requested handler.
    18  type StreamingRpcAck struct {
    19  	// Error is used to return whether an error occurred establishing the
    20  	// streaming RPC. This error occurs before entering the RPC handler.
    21  	Error string
    22  }
    23  
    24  // StreamingRpcHandler defines the handler for a streaming RPC.
    25  type StreamingRpcHandler func(conn io.ReadWriteCloser)
    26  
    27  // StreamingRpcRegistry is used to add and retrieve handlers
    28  type StreamingRpcRegistry struct {
    29  	registry map[string]StreamingRpcHandler
    30  }
    31  
    32  // NewStreamingRpcRegistry creates a new registry. All registrations of
    33  // handlers should be done before retrieving handlers.
    34  func NewStreamingRpcRegistry() *StreamingRpcRegistry {
    35  	return &StreamingRpcRegistry{
    36  		registry: make(map[string]StreamingRpcHandler),
    37  	}
    38  }
    39  
    40  // Register registers a new handler for the given method name
    41  func (s *StreamingRpcRegistry) Register(method string, handler StreamingRpcHandler) {
    42  	s.registry[method] = handler
    43  }
    44  
    45  // GetHandler returns a handler for the given method or an error if it doesn't exist.
    46  func (s *StreamingRpcRegistry) GetHandler(method string) (StreamingRpcHandler, error) {
    47  	h, ok := s.registry[method]
    48  	if !ok {
    49  		return nil, fmt.Errorf("%s: %q", ErrUnknownMethod, method)
    50  	}
    51  
    52  	return h, nil
    53  }
    54  
    55  // Bridge is used to just link two connections together and copy traffic
    56  func Bridge(a, b io.ReadWriteCloser) {
    57  	wg := sync.WaitGroup{}
    58  	wg.Add(2)
    59  	go func() {
    60  		defer wg.Done()
    61  		io.Copy(a, b)
    62  		a.Close()
    63  		b.Close()
    64  	}()
    65  	go func() {
    66  		defer wg.Done()
    67  		io.Copy(b, a)
    68  		a.Close()
    69  		b.Close()
    70  	}()
    71  	wg.Wait()
    72  }