From 2141e9631c55443a4b42bb721d028276e6e63394 Mon Sep 17 00:00:00 2001
From: Manish <itzmanish108@gmail.com>
Date: Sat, 7 Nov 2020 14:55:33 +0530
Subject: [PATCH] generate SendAndClose() and CloseAndRecv() on streams.

---
 cmd/protoc-gen-micro/plugin/micro/micro.go | 32 ++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/cmd/protoc-gen-micro/plugin/micro/micro.go b/cmd/protoc-gen-micro/plugin/micro/micro.go
index cb7d429e..c65af8d0 100644
--- a/cmd/protoc-gen-micro/plugin/micro/micro.go
+++ b/cmd/protoc-gen-micro/plugin/micro/micro.go
@@ -378,6 +378,11 @@ func (g *micro) generateClientMethod(reqServ, servName, serviceDescVar string, m
 	g.P("Context() context.Context")
 	g.P("SendMsg(interface{}) error")
 	g.P("RecvMsg(interface{}) error")
+
+	if genSend && !genRecv {
+		// client streaming, the server will send a response upon close
+		g.P("CloseAndRecv() (*", outType, ", error)")
+	}
 	g.P("Close() error")
 
 	if genSend {
@@ -394,6 +399,18 @@ func (g *micro) generateClientMethod(reqServ, servName, serviceDescVar string, m
 	g.P("}")
 	g.P()
 
+	if genSend && !genRecv {
+		// client streaming, the server will send a response upon close
+		g.P("func (x *", streamType, ") CloseAndRecv() (*", outType, ", error) {")
+		g.P("if err := x.stream.Close(); err != nil {")
+		g.P("return nil, err")
+		g.P("}")
+		g.P("r := new(", outType, ")")
+		g.P("err := x.RecvMsg(r)")
+		g.P("return r, err")
+		g.P("}")
+		g.P()
+	}
 	g.P("func (x *", streamType, ") Close() error {")
 	g.P("return x.stream.Close()")
 	g.P("}")
@@ -493,6 +510,10 @@ func (g *micro) generateServerMethod(servName string, method *pb.MethodDescripto
 	g.P("Context() context.Context")
 	g.P("SendMsg(interface{}) error")
 	g.P("RecvMsg(interface{}) error")
+	if !genSend {
+		// client streaming, the server will send a response upon close
+		g.P("SendAndClose(*", outType, ")  error")
+	}
 	g.P("Close() error")
 
 	if genSend {
@@ -511,6 +532,17 @@ func (g *micro) generateServerMethod(servName string, method *pb.MethodDescripto
 	g.P("}")
 	g.P()
 
+	if !genSend {
+		// client streaming, the server will send a response upon close
+		g.P("func (x *", streamType, ") SendAndClose(in *", outType, ") error {")
+		g.P("if err := x.SendMsg(in); err != nil {")
+		g.P("return err")
+		g.P("}")
+		g.P("return x.stream.Close()")
+		g.P("}")
+		g.P()
+	}
+	// other types of rpc don't send a response when the stream closes
 	g.P("func (x *", streamType, ") Close() error {")
 	g.P("return x.stream.Close()")
 	g.P("}")