minor changes

Author: YourTechBud

utils/client/client.go

@@ -22,7 +22,7 @@ type Client interface {
 }
 
 // DataCallback is the callback invoked when data is read by the socket
-type DataCallback func(data *model.Message)
+type DataCallback func(data *model.Message) bool
 
 // CreateWebsocketClient makes a client object to manage the socket
 func CreateWebsocketClient(socket *websocket.Conn) *WebsocketClient {

utils/client/grpc_realtime.go

@@ -93,7 +93,12 @@ func (c *GRPCRealtimeClient) Read(cb DataCallback) {
 		data["Where"] = temp
 
 		msg := &model.Message{Type: in.Type, ID: in.Id, Data: data}
-		cb(msg)
+
+		// Close the reader if callback returned false
+		next := cb(msg)
+		if !next {
+			return
+		}
 	}
 }
 

utils/client/grpc_service.go

@@ -89,7 +89,12 @@ func (c *GRPCServiceClient) Read(cb DataCallback) {
 		case utils.TypeServiceRegister:
 			data := map[string]interface{}{"service": in.Service, "token": in.Token, "project": in.Project}
 			msg := &model.Message{ID: in.Id, Type: utils.TypeServiceRegister, Data: data}
-			cb(msg)
+
+			// Close the reader if callback returned false
+			next := cb(msg)
+			if !next {
+				return
+			}
 
 		case utils.TypeServiceRequest:
 			var params interface{}
@@ -100,7 +105,12 @@ func (c *GRPCServiceClient) Read(cb DataCallback) {
 				"error":  in.Error,
 			}
 			msg := &model.Message{ID: in.Id, Type: utils.TypeServiceRequest, Data: data}
-			cb(msg)
+
+			// Close the reader if callback returned false
+			next := cb(msg)
+			if !next {
+				return
+			}
 
 		default:
 			log.Println("GRPC Service Error - Invalid request type", in.Type)

utils/client/websocket.go

@@ -54,7 +54,11 @@ func (c *WebsocketClient) Read(cb DataCallback) {
 			return
 		}
 
-		cb(data)
+		// Close the reader if callback returned false
+		next := cb(data)
+		if !next {
+			return
+		}
 	}
 }
 

utils/server/grpc.go

@@ -355,7 +355,7 @@ func (s *Server) Service(stream pb.SpaceCloud_ServiceServer) error {
 	// Get GRPC Service client details
 	clientID := c.ClientID()
 
-	c.Read(func(req *model.Message) {
+	c.Read(func(req *model.Message) bool {
 		switch req.Type {
 		case utils.TypeServiceRegister:
 			// TODO add security rule for functions registered as well
@@ -374,6 +374,8 @@ func (s *Server) Service(stream pb.SpaceCloud_ServiceServer) error {
 
 			s.functions.HandleServiceResponse(data)
 		}
+
+		return true
 	})
 	return nil
 }
@@ -388,7 +390,7 @@ func (s *Server) RealTime(stream pb.SpaceCloud_RealTimeServer) error {
 	ctx := c.Context()
 	clientID := c.ClientID()
 
-	c.Read(func(req *model.Message) {
+	c.Read(func(req *model.Message) bool {
 		switch req.Type {
 		case utils.TypeRealtimeSubscribe:
 			// For realtime subscribe event
@@ -402,7 +404,7 @@ func (s *Server) RealTime(stream pb.SpaceCloud_RealTimeServer) error {
 			if err != nil {
 				res := model.RealtimeResponse{Group: data.Group, ID: data.ID, Ack: false, Error: err.Error()}
 				c.Write(&model.Message{ID: req.ID, Type: req.Type, Data: res})
-				return
+				return true
 			}
 
 			// Send response to c
@@ -420,6 +422,8 @@ func (s *Server) RealTime(stream pb.SpaceCloud_RealTimeServer) error {
 			res := model.RealtimeResponse{Group: data.Group, ID: data.ID, Ack: true}
 			c.Write(&model.Message{ID: req.ID, Type: req.Type, Data: res})
 		}
+
+		return true
 	})
 
 	return nil

utils/server/websocket.go

@@ -37,7 +37,7 @@ func (s *Server) handleWebsocket() http.HandlerFunc {
 		ctx := c.Context()
 		clientID := c.ClientID()
 
-		c.Read(func(req *model.Message) {
+		c.Read(func(req *model.Message) bool {
 			switch req.Type {
 			case utils.TypeRealtimeSubscribe:
 				// For realtime subscribe event
@@ -51,7 +51,7 @@ func (s *Server) handleWebsocket() http.HandlerFunc {
 				if err != nil {
 					res := model.RealtimeResponse{Group: data.Group, ID: data.ID, Ack: false, Error: err.Error()}
 					c.Write(&model.Message{ID: req.ID, Type: req.Type, Data: res})
-					return
+					return true
 				}
 
 				// Send response to c
@@ -86,6 +86,7 @@ func (s *Server) handleWebsocket() http.HandlerFunc {
 
 				s.functions.HandleServiceResponse(data)
 			}
+			return true
 		})
 	}
 }