cobra

pkg/common/cmd/cron_task.go

@@ -10,8 +10,13 @@ func NewCronTaskCmd() *CronTaskCmd {
 	return &CronTaskCmd{NewRootCmd()}
 }
 
-func (c *CronTaskCmd) AddRunE(f func() error) {
+func (c *CronTaskCmd) addRunE(f func() error) {
 	c.Command.RunE = func(cmd *cobra.Command, args []string) error {
 		return f()
 	}
 }
+
+func (c *CronTaskCmd) Exec(f func() error) error {
+	c.addRunE(f)
+	return c.Execute()
+}

pkg/common/cmd/msg_gateway.go

@@ -1 +1,36 @@
 package cmd
+
+import (
+	"OpenIM/internal/msggateway"
+	"OpenIM/pkg/common/constant"
+	"github.com/spf13/cobra"
+)
+
+type MsgGatewayCmd struct {
+	*RootCmd
+}
+
+func NewMsgGatewayCmd() MsgGatewayCmd {
+	return MsgGatewayCmd{NewRootCmd()}
+}
+
+func (m *MsgGatewayCmd) AddWsPortFlag() {
+	m.Command.Flags().IntP(constant.FlagWsPort, "w", 0, "ws server listen port")
+}
+
+func (m *MsgGatewayCmd) getWsPortFlag(cmd *cobra.Command) int {
+	port, _ := cmd.Flags().GetInt(constant.FlagWsPort)
+	return port
+}
+
+func (m *MsgGatewayCmd) addRun() {
+	m.Command.Run = func(cmd *cobra.Command, args []string) {
+		msggateway.Init(m.getPortFlag(cmd), m.getWsPortFlag(cmd))
+		msggateway.Run(m.getPrometheusPortFlag(cmd))
+	}
+}
+
+func (m *MsgGatewayCmd) Exec() error {
+	m.addRun()
+	return m.Execute()
+}

pkg/common/cmd/msg_transfer.go

@@ -0,0 +1,25 @@
+package cmd
+
+import (
+	"OpenIM/internal/msgtransfer"
+	"github.com/spf13/cobra"
+)
+
+type MsgTransferCmd struct {
+	*RootCmd
+}
+
+func NewMsgTransferCmd() MsgTransferCmd {
+	return MsgTransferCmd{NewRootCmd()}
+}
+
+func (m *MsgTransferCmd) addRunE() {
+	m.Command.RunE = func(cmd *cobra.Command, args []string) error {
+		return msgtransfer.StartTransfer(m.getPrometheusPortFlag(cmd))
+	}
+}
+
+func (m *MsgTransferCmd) Exec() error {
+	m.addRunE()
+	return m.Execute()
+}

pkg/common/cmd/push.go

@@ -0,0 +1,22 @@
+package cmd
+
+import (
+	"OpenIM/internal/push"
+	"OpenIM/internal/startrpc"
+	"OpenIM/pkg/common/config"
+	"github.com/spf13/cobra"
+)
+
+type PushCmd struct {
+	*RpcCmd
+}
+
+func NewPushCmd() *PushCmd {
+	return &PushCmd{NewRpcCmd(config.Config.RpcRegisterName.OpenImPushName)}
+}
+
+func (r *RpcCmd) AddPush() {
+	r.Command.RunE = func(cmd *cobra.Command, args []string) error {
+		return startrpc.Start(r.getPortFlag(cmd), r.rpcRegisterName, r.getPrometheusPortFlag(cmd), push.Start)
+	}
+}

pkg/common/cmd/root.go

@@ -8,9 +8,6 @@ import (
 
 type RootCmd struct {
 	Command cobra.Command
-	port    int
-
-	prometheusPort int
 }
 
 func NewRootCmd() (rootCmd *RootCmd) {
@@ -20,8 +17,6 @@ func NewRootCmd() (rootCmd *RootCmd) {
 		Short: "Start the server",
 		Long:  `Start the server`,
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
-			rootCmd.port = rootCmd.getPortFlag(cmd)
-			rootCmd.prometheusPort = rootCmd.getPrometheusPortFlag(cmd)
 			return rootCmd.getConfFromCmdAndInit(cmd)
 		},
 	}
@@ -38,7 +33,7 @@ func (r *RootCmd) AddRunE(f func(cmd RootCmd) error) {
 
 func (r *RootCmd) AddRpc(f func(port, prometheusPort int) error) {
 	r.Command.RunE = func(cmd *cobra.Command, args []string) error {
-		return f(r.port, r.prometheusPort)
+		return f(r.getPortFlag(cmd), r.getPrometheusPortFlag(cmd))
 	}
 }
 
@@ -55,23 +50,15 @@ func (r *RootCmd) getPortFlag(cmd *cobra.Command) int {
 	return port
 }
 
-func (r *RootCmd) GetPortFlag() int {
-	return r.port
-}
-
 func (r *RootCmd) AddPrometheusPortFlag() {
-	r.Command.Flags().StringP(constant.PrometheusPort, "pp", "", "server listen port")
+	r.Command.Flags().StringP(constant.FlagPrometheusPort, "pp", "", "server prometheus listen port")
 }
 
 func (r *RootCmd) getPrometheusPortFlag(cmd *cobra.Command) int {
-	port, _ := cmd.Flags().GetInt(constant.PrometheusPort)
+	port, _ := cmd.Flags().GetInt(constant.FlagPrometheusPort)
 	return port
 }
 
-func (r *RootCmd) GetPrometheusPortFlag() int {
-	return r.prometheusPort
-}
-
 func (r *RootCmd) getConfFromCmdAndInit(cmdLines *cobra.Command) error {
 	configFolderPath, _ := cmdLines.Flags().GetString(constant.FlagConf)
 	return config.InitConfig(configFolderPath)

pkg/common/cmd/rpc.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"OpenIM/internal/startrpc"
 	"OpenIM/pkg/discoveryregistry"
 	"github.com/spf13/cobra"
 	"google.golang.org/grpc"
@@ -8,14 +9,20 @@ import (
 
 type RpcCmd struct {
 	*RootCmd
+	rpcRegisterName string
 }
 
-func NewRpcCmd() *RpcCmd {
-	return &RpcCmd{NewRootCmd()}
+func NewRpcCmd(rpcRegisterName string) *RpcCmd {
+	return &RpcCmd{NewRootCmd(), rpcRegisterName}
 }
 
-func (r *RpcCmd) AddRpc(f func(port, rpcRegisterName string, prometheusPort int, rpcFn func(client discoveryregistry.SvcDiscoveryRegistry, server *grpc.Server) error, options ...grpc.ServerOption) error) {
+func (r *RpcCmd) AddRpc(rpcFn func(client discoveryregistry.SvcDiscoveryRegistry, server *grpc.Server) error) {
 	r.Command.RunE = func(cmd *cobra.Command, args []string) error {
-		return f(r.port, r.prometheusPort)
+		return startrpc.Start(r.getPortFlag(cmd), r.rpcRegisterName, r.getPrometheusPortFlag(cmd), rpcFn)
 	}
 }
+
+func (r *RpcCmd) Exec(rpcFn func(client discoveryregistry.SvcDiscoveryRegistry, server *grpc.Server) error) error {
+	r.AddRpc(rpcFn)
+	return r.Execute()
+}