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