github.com/homburg/packer@v0.6.1-0.20140528012651-1dcaf1716848/packer/rpc/server.go (about)

     1  package rpc
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/mitchellh/packer/packer"
     6  	"github.com/ugorji/go/codec"
     7  	"io"
     8  	"log"
     9  	"net/rpc"
    10  	"sync/atomic"
    11  )
    12  
    13  var endpointId uint64
    14  
    15  const (
    16  	DefaultArtifactEndpoint      string = "Artifact"
    17  	DefaultBuildEndpoint                = "Build"
    18  	DefaultBuilderEndpoint              = "Builder"
    19  	DefaultCacheEndpoint                = "Cache"
    20  	DefaultCommandEndpoint              = "Command"
    21  	DefaultCommunicatorEndpoint         = "Communicator"
    22  	DefaultEnvironmentEndpoint          = "Environment"
    23  	DefaultHookEndpoint                 = "Hook"
    24  	DefaultPostProcessorEndpoint        = "PostProcessor"
    25  	DefaultProvisionerEndpoint          = "Provisioner"
    26  	DefaultUiEndpoint                   = "Ui"
    27  )
    28  
    29  // Server represents an RPC server for Packer. This must be paired on
    30  // the other side with a Client.
    31  type Server struct {
    32  	mux      *MuxConn
    33  	streamId uint32
    34  	server   *rpc.Server
    35  	closeMux bool
    36  }
    37  
    38  // NewServer returns a new Packer RPC server.
    39  func NewServer(conn io.ReadWriteCloser) *Server {
    40  	result := newServerWithMux(NewMuxConn(conn), 0)
    41  	result.closeMux = true
    42  	return result
    43  }
    44  
    45  func newServerWithMux(mux *MuxConn, streamId uint32) *Server {
    46  	return &Server{
    47  		mux:      mux,
    48  		streamId: streamId,
    49  		server:   rpc.NewServer(),
    50  		closeMux: false,
    51  	}
    52  }
    53  
    54  func (s *Server) Close() error {
    55  	if s.closeMux {
    56  		log.Printf("[WARN] Shutting down mux conn in Server")
    57  		return s.mux.Close()
    58  	}
    59  
    60  	return nil
    61  }
    62  
    63  func (s *Server) RegisterArtifact(a packer.Artifact) {
    64  	s.server.RegisterName(DefaultArtifactEndpoint, &ArtifactServer{
    65  		artifact: a,
    66  	})
    67  }
    68  
    69  func (s *Server) RegisterBuild(b packer.Build) {
    70  	s.server.RegisterName(DefaultBuildEndpoint, &BuildServer{
    71  		build: b,
    72  		mux:   s.mux,
    73  	})
    74  }
    75  
    76  func (s *Server) RegisterBuilder(b packer.Builder) {
    77  	s.server.RegisterName(DefaultBuilderEndpoint, &BuilderServer{
    78  		builder: b,
    79  		mux:     s.mux,
    80  	})
    81  }
    82  
    83  func (s *Server) RegisterCache(c packer.Cache) {
    84  	s.server.RegisterName(DefaultCacheEndpoint, &CacheServer{
    85  		cache: c,
    86  	})
    87  }
    88  
    89  func (s *Server) RegisterCommand(c packer.Command) {
    90  	s.server.RegisterName(DefaultCommandEndpoint, &CommandServer{
    91  		command: c,
    92  		mux:     s.mux,
    93  	})
    94  }
    95  
    96  func (s *Server) RegisterCommunicator(c packer.Communicator) {
    97  	s.server.RegisterName(DefaultCommunicatorEndpoint, &CommunicatorServer{
    98  		c:   c,
    99  		mux: s.mux,
   100  	})
   101  }
   102  
   103  func (s *Server) RegisterEnvironment(b packer.Environment) {
   104  	s.server.RegisterName(DefaultEnvironmentEndpoint, &EnvironmentServer{
   105  		env: b,
   106  		mux: s.mux,
   107  	})
   108  }
   109  
   110  func (s *Server) RegisterHook(h packer.Hook) {
   111  	s.server.RegisterName(DefaultHookEndpoint, &HookServer{
   112  		hook: h,
   113  		mux:  s.mux,
   114  	})
   115  }
   116  
   117  func (s *Server) RegisterPostProcessor(p packer.PostProcessor) {
   118  	s.server.RegisterName(DefaultPostProcessorEndpoint, &PostProcessorServer{
   119  		mux: s.mux,
   120  		p:   p,
   121  	})
   122  }
   123  
   124  func (s *Server) RegisterProvisioner(p packer.Provisioner) {
   125  	s.server.RegisterName(DefaultProvisionerEndpoint, &ProvisionerServer{
   126  		mux: s.mux,
   127  		p:   p,
   128  	})
   129  }
   130  
   131  func (s *Server) RegisterUi(ui packer.Ui) {
   132  	s.server.RegisterName(DefaultUiEndpoint, &UiServer{
   133  		ui: ui,
   134  	})
   135  }
   136  
   137  // ServeConn serves a single connection over the RPC server. It is up
   138  // to the caller to obtain a proper io.ReadWriteCloser.
   139  func (s *Server) Serve() {
   140  	// Accept a connection on stream ID 0, which is always used for
   141  	// normal client to server connections.
   142  	stream, err := s.mux.Accept(s.streamId)
   143  	defer stream.Close()
   144  	if err != nil {
   145  		log.Printf("[ERR] Error retrieving stream for serving: %s", err)
   146  		return
   147  	}
   148  
   149  	var h codec.MsgpackHandle
   150  	rpcCodec := codec.GoRpc.ServerCodec(stream, &h)
   151  	s.server.ServeCodec(rpcCodec)
   152  }
   153  
   154  // registerComponent registers a single Packer RPC component onto
   155  // the RPC server. If id is true, then a unique ID number will be appended
   156  // onto the end of the endpoint.
   157  //
   158  // The endpoint name is returned.
   159  func registerComponent(server *rpc.Server, name string, rcvr interface{}, id bool) string {
   160  	endpoint := name
   161  	if id {
   162  		fmt.Sprintf("%s.%d", endpoint, atomic.AddUint64(&endpointId, 1))
   163  	}
   164  
   165  	server.RegisterName(endpoint, rcvr)
   166  	return endpoint
   167  }