github.com/kaixiang/packer@v0.5.2-0.20140114230416-1f5786b0d7f1/packer/rpc/server.go (about)

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