github.com/rahart/packer@v0.12.2-0.20161229105310-282bb6ad370f/packer/rpc/server.go (about)

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