diff --git a/cmd/client.go b/cmd/client.go index 4e7e3e1..ae89930 100644 --- a/cmd/client.go +++ b/cmd/client.go @@ -29,7 +29,7 @@ func ClientCmd() *cobra.Command { // clientCmd represents the client command clientCmd := &cobra.Command{ Use: "client", - Short: "A handy tool that can call the Squeeze's API.", + Short: "A handy tools that can call the Squeeze's API.", Long: `This command allows you to interact with Squeeze and stress targets with multiple protocols. Currently supported protocol is only http, other protocols are under development. Look forward to your contribution. diff --git a/cmd/slave.go b/cmd/slave.go index 07b5f87..a8e8c9b 100644 --- a/cmd/slave.go +++ b/cmd/slave.go @@ -55,19 +55,19 @@ func SlaveCmd() *cobra.Command { }, } - slaveCmd.PersistentFlags().StringVar(&serverArgs.HTTPAddr, "httpAddr", ":9998", + slaveCmd.PersistentFlags().StringVar(&serverArgs.HTTPAddr, "httpAddr", ":9996", "Squeeze service HTTP address") - slaveCmd.PersistentFlags().StringVar(&serverArgs.GRPCAddr, "grpcAddr", ":9997", + slaveCmd.PersistentFlags().StringVar(&serverArgs.GRPCAddr, "grpcAddr", ":9995", "Squeeze service grpc address") - slaveCmd.PersistentFlags().StringVar(&serverArgs.MasterAddr, "masterAddr", "", - "The address of the master server") + slaveCmd.PersistentFlags().StringVar(&serverArgs.HttpMasterAddr, "httpMasterAddr", "", + "Master Server's http address.") slaveCmd.PersistentFlags().StringVar(&serverArgs.GrpcMasterAddr, "grpcMasterAddr", "", - "The address of the grpc master server") + "Master Server's grpc address.") slaveCmd.PersistentFlags().DurationVar(&serverArgs.ReportInterval, "reportInterval", 5, "Task reporting interval to the master") slaveCmd.PersistentFlags().IntVar(&serverArgs.ResultCapacity, "resultCapacity", 2000, "The capacity of the results channel for aggregating.") - slaveCmd.MarkPersistentFlagRequired("masterAddr") + slaveCmd.MarkPersistentFlagRequired("httpMasterAddr") slaveCmd.MarkPersistentFlagRequired("grpcMasterAddr") return slaveCmd diff --git a/cmd/web.go b/cmd/web.go index 7daa699..9b1c5f9 100644 --- a/cmd/web.go +++ b/cmd/web.go @@ -62,8 +62,8 @@ func WebCmd() *cobra.Command { webCmd.PersistentFlags().StringVar(&serverArgs.HTTPAddr, "httpAddr", ":9991", "The address and port of the web server.") - webCmd.PersistentFlags().StringVar(&serverArgs.MasterAddr, "masterAddr", "", - "The address of the master server") + webCmd.PersistentFlags().StringVar(&serverArgs.HttpMasterAddr, "masterAddr", "", + "Master server's http address.") webCmd.PersistentFlags().StringVar(&webOptions.DSN, "dsn", "", `Data Source Name. If you specify --type=mysql, need to set this option. Format: username:password@protocol(address)/dbname?param=value`) diff --git a/pkg/server/proxy.go b/pkg/server/proxy.go index ca601a5..b10e9c0 100644 --- a/pkg/server/proxy.go +++ b/pkg/server/proxy.go @@ -15,10 +15,11 @@ package server import ( - "github.com/agile6v/squeeze/pkg/version" + "fmt" + "net/url" "net/http" "net/http/httputil" - "net/url" + "github.com/agile6v/squeeze/pkg/version" ) type Proxy struct { @@ -26,10 +27,13 @@ type Proxy struct { proxy *httputil.ReverseProxy } -func NewProxy(target string) *Proxy { - url, _ := url.Parse(target) +func NewProxy(target string) (*Proxy, error) { + url, err := url.Parse(fmt.Sprintf("http://%v/", target)) + if err != nil { + return nil, err + } - return &Proxy{target: url, proxy: httputil.NewSingleHostReverseProxy(url)} + return &Proxy{target: url, proxy: httputil.NewSingleHostReverseProxy(url)}, nil } func (p *Proxy) handle(w http.ResponseWriter, r *http.Request) { diff --git a/pkg/server/server.go b/pkg/server/server.go index 1b39943..931c618 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -65,7 +65,7 @@ type ServerArgs struct { HTTPAddr string // The listening address for http GRPCAddr string // The listening address for grpc - MasterAddr string // Master's HTTP Address + HttpMasterAddr string // Master's HTTP Address GrpcMasterAddr string // Master's GRPC Address ReportInterval time.Duration // Heartbeat reporting interval ResultCapacity int diff --git a/pkg/server/slave_server.go b/pkg/server/slave_server.go index 75a5daf..4de54f5 100644 --- a/pkg/server/slave_server.go +++ b/pkg/server/slave_server.go @@ -27,6 +27,7 @@ import ( "github.com/agile6v/squeeze/pkg/pb" log "github.com/golang/glog" "github.com/agile6v/squeeze/pkg/proto/builder" + "github.com/agile6v/squeeze/pkg/util" ) type SlaveServer struct { @@ -40,13 +41,28 @@ func (s *SlaveServer) Initialize(args *ServerArgs) error { s.ServerBase.Initialize(args) s.Mode = Slave + // check if the HttpMasterAddr is valid + _, _, err := util.GetHostPort(s.args.HttpMasterAddr) + if err != nil { + return err + } + + // check if the GrpcMasterAddr is valid + _, _, err = util.GetHostPort(s.args.GrpcMasterAddr) + if err != nil { + return err + } + // Create the gRPC server // TODO: create the grpc options s.grpcServer = grpc.NewServer(grpc.MaxConcurrentStreams(256), grpc.MaxMsgSize(1024*1024)) pb.RegisterSqueezeServiceServer(s.grpcServer, s) - // TODO: check if the MasterAddr is valid - proxy := NewProxy(s.args.MasterAddr) + proxy, err := NewProxy(s.args.HttpMasterAddr) + if err != nil { + return err + } + http.HandleFunc("/", proxy.handle) return nil } diff --git a/pkg/server/web_server.go b/pkg/server/web_server.go index 38129e6..31f699d 100644 --- a/pkg/server/web_server.go +++ b/pkg/server/web_server.go @@ -43,7 +43,7 @@ func (s *WebServer) Initialize(args *ServerArgs) error { } api := &api.AppAPI{ - MasterAddr: s.args.MasterAddr, + MasterAddr: s.args.HttpMasterAddr, HTTPAddr: s.args.HTTPAddr, LocalAddr: ip + ":" + port, }