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  }