[ISSIUE #75] support shutdown mehtod

internal/client.go

@@ -43,7 +43,7 @@ const (
 	// Pulling topic information interval from the named server
 	_PullNameServerInterval = 30 * time.Second
 
-	// Pulling topic information interval from the named server
+	// Sending heart beat interval to all broker
 	_HeartbeatBrokerInterval = 30 * time.Second
 
 	// Offset persistent interval for consumer
@@ -54,7 +54,7 @@ const (
 )
 
 var (
-	ErrServiceState = errors.New("service state is not running, please check")
+	ErrServiceState = errors.New("service close is not running, please check")
 
 	_VIPChannelEnable = false
 )
@@ -129,6 +129,7 @@ type RMQClient struct {
 
 	remoteClient *remote.RemotingClient
 	hbMutex      sync.Mutex
+	close        bool
 }
 
 var clientMap sync.Map
@@ -150,6 +151,9 @@ func GetOrNewRocketMQClient(option ClientOptions) *RMQClient {
 }
 
 func (c *RMQClient) Start() {
+	//ctx, cancel := context.WithCancel(context.Background())
+	//c.cancel = cancel
+	c.close = false
 	c.once.Do(func() {
 		// TODO fetchNameServerAddr
 		go func() {}()
@@ -158,15 +162,15 @@ func (c *RMQClient) Start() {
 		go func() {
 			// delay
 			time.Sleep(50 * time.Millisecond)
-			for {
+			for !c.close{
 				c.UpdateTopicRouteInfo()
 				time.Sleep(_PullNameServerInterval)
 			}
 		}()
 
 		// TODO cleanOfflineBroker & sendHeartbeatToAllBrokerWithLock
 		go func() {
-			for {
+			for !c.close{
 				cleanOfflineBroker()
 				c.SendHeartbeatToAllBrokerWithLock()
 				time.Sleep(_HeartbeatBrokerInterval)
@@ -176,7 +180,7 @@ func (c *RMQClient) Start() {
 		// schedule persist offset
 		go func() {
 			//time.Sleep(10 * time.Second)
-			for {
+			for !c.close{
 				c.consumerMap.Range(func(key, value interface{}) bool {
 					consumer := value.(InnerConsumer)
 					consumer.PersistConsumerOffset()
@@ -187,7 +191,7 @@ func (c *RMQClient) Start() {
 		}()
 
 		go func() {
-			for {
+			for !c.close{
 				c.RebalanceImmediately()
 				time.Sleep(_RebalanceInterval)
 			}
@@ -196,7 +200,8 @@ func (c *RMQClient) Start() {
 }
 
 func (c *RMQClient) Shutdown() {
-	// TODO
+	c.remoteClient.ShutDown()
+	c.close = true
 }
 
 func (c *RMQClient) ClientID() string {
@@ -209,18 +214,28 @@ func (c *RMQClient) ClientID() string {
 
 func (c *RMQClient) InvokeSync(addr string, request *remote.RemotingCommand,
 	timeoutMillis time.Duration) (*remote.RemotingCommand, error) {
+	if c.close {
+		return nil, ErrServiceState
+	}
 	return c.remoteClient.InvokeSync(addr, request, timeoutMillis)
 }
 
 func (c *RMQClient) InvokeAsync(addr string, request *remote.RemotingCommand,
 	timeoutMillis time.Duration, f func(*remote.RemotingCommand, error)) error {
+	if c.close {
+		return ErrServiceState
+	}
 	return c.remoteClient.InvokeAsync(addr, request, timeoutMillis, func(future *remote.ResponseFuture) {
 		f(future.ResponseCommand, future.Err)
 	})
+
 }
 
 func (c *RMQClient) InvokeOneWay(addr string, request *remote.RemotingCommand,
 	timeoutMillis time.Duration) error {
+	if c.close {
+		return ErrServiceState
+	}
 	return c.remoteClient.InvokeOneWay(addr, request, timeoutMillis)
 }
 

internal/remote/codec.go

@@ -73,7 +73,7 @@ func NewRemotingCommand(code int16, header CustomHeader, body []byte) *RemotingC
 }
 
 func (command *RemotingCommand) String() string {
-	return fmt.Sprintf("Code: %d, Opaque: %d, Remark: %s, ExtFields: %v",
+	return fmt.Sprintf("Code: %d, opaque: %d, Remark: %s, ExtFields: %v",
 		command.Code, command.Opaque, command.Remark, command.ExtFields)
 }
 

internal/remote/codec_test.go

@@ -96,7 +96,7 @@ func Test_decode(t *testing.T) {
 			t.Fatalf("wrong Version. want=%d, got=%d", rc.Version, decodedRc.Version)
 		}
 		if rc.Opaque != decodedRc.Opaque {
-			t.Fatalf("wrong Opaque. want=%d, got=%d", rc.Opaque, decodedRc.Opaque)
+			t.Fatalf("wrong opaque. want=%d, got=%d", rc.Opaque, decodedRc.Opaque)
 		}
 		if rc.Remark != decodedRc.Remark {
 			t.Fatalf("wrong remark. want=%s, got=%s", rc.Remark, decodedRc.Remark)
@@ -167,7 +167,7 @@ func Test_jsonCodec_decodeHeader(t *testing.T) {
 			t.Fatalf("wrong Version. want=%d, got=%d", rc.Version, decodedRc.Version)
 		}
 		if rc.Opaque != decodedRc.Opaque {
-			t.Fatalf("wrong Opaque. want=%d, got=%d", rc.Opaque, decodedRc.Opaque)
+			t.Fatalf("wrong opaque. want=%d, got=%d", rc.Opaque, decodedRc.Opaque)
 		}
 		if rc.Remark != decodedRc.Remark {
 			t.Fatalf("wrong remark. want=%s, got=%s", rc.Remark, decodedRc.Remark)
@@ -237,7 +237,7 @@ func Test_rmqCodec_decodeHeader(t *testing.T) {
 			t.Fatalf("wrong Version. want=%d, got=%d", rc.Version, decodedRc.Version)
 		}
 		if rc.Opaque != decodedRc.Opaque {
-			t.Fatalf("wrong Opaque. want=%d, got=%d", rc.Opaque, decodedRc.Opaque)
+			t.Fatalf("wrong opaque. want=%d, got=%d", rc.Opaque, decodedRc.Opaque)
 		}
 		if rc.Remark != decodedRc.Remark {
 			t.Fatalf("wrong remark. want=%s, got=%s", rc.Remark, decodedRc.Remark)

internal/remote/future.go

@@ -28,21 +28,21 @@ type ResponseFuture struct {
 	SendRequestOK   bool
 	Err             error
 	Opaque          int32
-	TimeoutMillis   time.Duration
+	Timeout         time.Duration
 	callback        func(*ResponseFuture)
-	BeginTimestamp  int64
+	BeginTimestamp  time.Duration
 	Done            chan bool
 	callbackOnce    sync.Once
 }
 
 // NewResponseFuture create ResponseFuture with opaque, timeout and callback
-func NewResponseFuture(opaque int32, timeoutMillis time.Duration, callback func(*ResponseFuture)) *ResponseFuture {
+func NewResponseFuture(opaque int32, timeout time.Duration, callback func(*ResponseFuture)) *ResponseFuture {
 	return &ResponseFuture{
 		Opaque:         opaque,
 		Done:           make(chan bool),
-		TimeoutMillis:  timeoutMillis,
+		Timeout:        timeout,
 		callback:       callback,
-		BeginTimestamp: time.Now().Unix() * 1000,
+		BeginTimestamp: time.Duration(time.Now().Unix()) * time.Second,
 	}
 }
 
@@ -55,16 +55,16 @@ func (r *ResponseFuture) executeInvokeCallback() {
 }
 
 func (r *ResponseFuture) isTimeout() bool {
-	diff := time.Now().Unix()*1000 - r.BeginTimestamp
-	return diff > int64(r.TimeoutMillis)
+	elapse := time.Duration(time.Now().Unix())*time.Second - r.BeginTimestamp
+	return elapse > r.Timeout
 }
 
 func (r *ResponseFuture) waitResponse() (*RemotingCommand, error) {
 	var (
 		cmd *RemotingCommand
 		err error
 	)
-	timer := time.NewTimer(r.TimeoutMillis * time.Millisecond)
+	timer := time.NewTimer(r.Timeout)
 	for {
 		select {
 		case <-r.Done:

internal/remote/remote_client.go

@@ -32,7 +32,6 @@ import (
 var (
 	//ErrRequestTimeout for request timeout error
 	ErrRequestTimeout = errors.New("request timeout")
-	connectionLocker  sync.Mutex
 )
 
 type ClientRequestFunc func(*RemotingCommand) *RemotingCommand
@@ -42,10 +41,11 @@ type TcpOption struct {
 }
 
 type RemotingClient struct {
-	responseTable   sync.Map
-	connectionTable sync.Map
-	option          TcpOption
-	processors      map[int16]ClientRequestFunc
+	responseTable    sync.Map
+	connectionTable  sync.Map
+	option           TcpOption
+	processors       map[int16]ClientRequestFunc
+	connectionLocker sync.Mutex
 }
 
 func NewRemotingClient() *RemotingClient {
@@ -59,29 +59,29 @@ func (c *RemotingClient) RegisterRequestFunc(code int16, f ClientRequestFunc) {
 }
 
 // TODO: merge sync and async model. sync should run on async model by blocking on chan
-func (c *RemotingClient) InvokeSync(addr string, request *RemotingCommand, timeoutMillis time.Duration) (*RemotingCommand, error) {
+func (c *RemotingClient) InvokeSync(addr string, request *RemotingCommand, timeout time.Duration) (*RemotingCommand, error) {
 	conn, err := c.connect(addr)
 	if err != nil {
 		return nil, err
 	}
-	resp := NewResponseFuture(request.Opaque, timeoutMillis, nil)
+	resp := NewResponseFuture(request.Opaque, timeout, nil)
 	c.responseTable.Store(resp.Opaque, resp)
-	err = c.sendRequest(conn, request)
 	defer c.responseTable.Delete(request.Opaque)
+	err = c.sendRequest(conn, request)
 	if err != nil {
 		return nil, err
 	}
 	resp.SendRequestOK = true
 	return resp.waitResponse()
 }
 
-// InvokeAsync send request witout blocking, just return immediately.
-func (c *RemotingClient) InvokeAsync(addr string, request *RemotingCommand, timeoutMillis time.Duration, callback func(*ResponseFuture)) error {
+// InvokeAsync send request without blocking, just return immediately.
+func (c *RemotingClient) InvokeAsync(addr string, request *RemotingCommand, timeout time.Duration, callback func(*ResponseFuture)) error {
 	conn, err := c.connect(addr)
 	if err != nil {
 		return err
 	}
-	resp := NewResponseFuture(request.Opaque, timeoutMillis, callback)
+	resp := NewResponseFuture(request.Opaque, timeout, callback)
 	c.responseTable.Store(resp.Opaque, resp)
 	err = c.sendRequest(conn, request)
 	if err != nil {
@@ -107,27 +107,10 @@ func (c *RemotingClient) InvokeOneWay(addr string, request *RemotingCommand, tim
 	return c.sendRequest(conn, request)
 }
 
-func (c *RemotingClient) ScanResponseTable() {
-	rfs := make([]*ResponseFuture, 0)
-	c.responseTable.Range(func(key, value interface{}) bool {
-		if resp, ok := value.(*ResponseFuture); ok {
-			if (resp.BeginTimestamp + int64(resp.TimeoutMillis) + 1000) <= time.Now().Unix()*1000 {
-				rfs = append(rfs, resp)
-				c.responseTable.Delete(key)
-			}
-		}
-		return true
-	})
-	for _, rf := range rfs {
-		rf.Err = ErrRequestTimeout
-		rf.executeInvokeCallback()
-	}
-}
-
 func (c *RemotingClient) connect(addr string) (net.Conn, error) {
 	//it needs additional locker.
-	connectionLocker.Lock()
-	defer connectionLocker.Unlock()
+	c.connectionLocker.Lock()
+	defer c.connectionLocker.Unlock()
 	conn, ok := c.connectionTable.Load(addr)
 	if ok {
 		return conn.(net.Conn), nil
@@ -181,7 +164,7 @@ func (c *RemotingClient) receiveResponse(r net.Conn) {
 		}
 	}
 	if scanner.Err() != nil {
-		rlog.Errorf("net: %s scanner exit, err: %s.", r.RemoteAddr().String(), scanner.Err())
+		rlog.Errorf("net: %s scanner exit, Err: %s.", r.RemoteAddr().String(), scanner.Err())
 	} else {
 		rlog.Infof("net: %s scanner exit.", r.RemoteAddr().String())
 	}
@@ -237,3 +220,15 @@ func (c *RemotingClient) closeConnection(toCloseConn net.Conn) {
 		}
 	})
 }
+
+func (c *RemotingClient) ShutDown() {
+	c.responseTable.Range(func(key, value interface{}) bool {
+		c.responseTable.Delete(key)
+		return true
+	})
+	c.connectionTable.Range(func(key, value interface{}) bool {
+		conn := value.(net.Conn)
+		conn.Close()
+		return true
+	})
+}