github.com/askholme/packer@v0.7.2-0.20140924152349-70d9566a6852/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      *muxBroker
    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  	mux, _ := newMuxBrokerServer(conn)
    41  	result := newServerWithMux(mux, 0)
    42  	result.closeMux = true
    43  	go mux.Run()
    44  	return result
    45  }
    46  
    47  func newServerWithMux(mux *muxBroker, streamId uint32) *Server {
    48  	return &Server{
    49  		mux:      mux,
    50  		streamId: streamId,
    51  		server:   rpc.NewServer(),
    52  		closeMux: false,
    53  	}
    54  }
    55  
    56  func (s *Server) Close() error {
    57  	if s.closeMux {
    58  		log.Printf("[WARN] Shutting down mux conn in Server")
    59  		return s.mux.Close()
    60  	}
    61  
    62  	return nil
    63  }
    64  
    65  func (s *Server) RegisterArtifact(a packer.Artifact) {
    66  	s.server.RegisterName(DefaultArtifactEndpoint, &ArtifactServer{
    67  		artifact: a,
    68  	})
    69  }
    70  
    71  func (s *Server) RegisterBuild(b packer.Build) {
    72  	s.server.RegisterName(DefaultBuildEndpoint, &BuildServer{
    73  		build: b,
    74  		mux:   s.mux,
    75  	})
    76  }
    77  
    78  func (s *Server) RegisterBuilder(b packer.Builder) {
    79  	s.server.RegisterName(DefaultBuilderEndpoint, &BuilderServer{
    80  		builder: b,
    81  		mux:     s.mux,
    82  	})
    83  }
    84  
    85  func (s *Server) RegisterCache(c packer.Cache) {
    86  	s.server.RegisterName(DefaultCacheEndpoint, &CacheServer{
    87  		cache: c,
    88  	})
    89  }
    90  
    91  func (s *Server) RegisterCommand(c packer.Command) {
    92  	s.server.RegisterName(DefaultCommandEndpoint, &CommandServer{
    93  		command: c,
    94  		mux:     s.mux,
    95  	})
    96  }
    97  
    98  func (s *Server) RegisterCommunicator(c packer.Communicator) {
    99  	s.server.RegisterName(DefaultCommunicatorEndpoint, &CommunicatorServer{
   100  		c:   c,
   101  		mux: s.mux,
   102  	})
   103  }
   104  
   105  func (s *Server) RegisterEnvironment(b packer.Environment) {
   106  	s.server.RegisterName(DefaultEnvironmentEndpoint, &EnvironmentServer{
   107  		env: b,
   108  		mux: s.mux,
   109  	})
   110  }
   111  
   112  func (s *Server) RegisterHook(h packer.Hook) {
   113  	s.server.RegisterName(DefaultHookEndpoint, &HookServer{
   114  		hook: h,
   115  		mux:  s.mux,
   116  	})
   117  }
   118  
   119  func (s *Server) RegisterPostProcessor(p packer.PostProcessor) {
   120  	s.server.RegisterName(DefaultPostProcessorEndpoint, &PostProcessorServer{
   121  		mux: s.mux,
   122  		p:   p,
   123  	})
   124  }
   125  
   126  func (s *Server) RegisterProvisioner(p packer.Provisioner) {
   127  	s.server.RegisterName(DefaultProvisionerEndpoint, &ProvisionerServer{
   128  		mux: s.mux,
   129  		p:   p,
   130  	})
   131  }
   132  
   133  func (s *Server) RegisterUi(ui packer.Ui) {
   134  	s.server.RegisterName(DefaultUiEndpoint, &UiServer{
   135  		ui: ui,
   136  	})
   137  }
   138  
   139  // ServeConn serves a single connection over the RPC server. It is up
   140  // to the caller to obtain a proper io.ReadWriteCloser.
   141  func (s *Server) Serve() {
   142  	// Accept a connection on stream ID 0, which is always used for
   143  	// normal client to server connections.
   144  	stream, err := s.mux.Accept(s.streamId)
   145  	if err != nil {
   146  		log.Printf("[ERR] Error retrieving stream for serving: %s", err)
   147  		return
   148  	}
   149  	defer stream.Close()
   150  
   151  	var h codec.MsgpackHandle
   152  	rpcCodec := codec.GoRpc.ServerCodec(stream, &h)
   153  	s.server.ServeCodec(rpcCodec)
   154  }
   155  
   156  // registerComponent registers a single Packer RPC component onto
   157  // the RPC server. If id is true, then a unique ID number will be appended
   158  // onto the end of the endpoint.
   159  //
   160  // The endpoint name is returned.
   161  func registerComponent(server *rpc.Server, name string, rcvr interface{}, id bool) string {
   162  	endpoint := name
   163  	if id {
   164  		fmt.Sprintf("%s.%d", endpoint, atomic.AddUint64(&endpointId, 1))
   165  	}
   166  
   167  	server.RegisterName(endpoint, rcvr)
   168  	return endpoint
   169  }