Making sure goroutines exit properly, adding stream context handling

loopholelabs · Sep 8, 2024 · 5f96233 · 5f96233
1 parent 2a7e437
commit 5f96233
Show file tree

Hide file tree

Showing 9 changed files with 112 additions and 79 deletions.
diff --git a/go.mod b/go.mod
@@ -3,7 +3,7 @@ module github.com/loopholelabs/frpc-go
 go 1.22
 
 require (
-	github.com/loopholelabs/frisbee-go v0.9.2
+	github.com/loopholelabs/frisbee-go v0.9.3-0.20240908010142-de266dc4d891
 	github.com/loopholelabs/logging v0.3.0
 	github.com/loopholelabs/polyglot/v2 v2.0.2
 	github.com/loopholelabs/testing v0.2.3

diff --git a/go.sum b/go.sum
@@ -8,8 +8,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/loopholelabs/common v0.4.9 h1:9MPUYlZZ/qx3Kt8LXgXxcSXthrM91od8026c4DlGpAU=
 github.com/loopholelabs/common v0.4.9/go.mod h1:Wop5srN1wYT+mdQ9gZ+kn2I9qKAyVd0FB48pThwIa9M=
-github.com/loopholelabs/frisbee-go v0.9.2 h1:TiihZMh/aq82olM3JEY0h5WiSHOji9lheaQ26V35zK0=
-github.com/loopholelabs/frisbee-go v0.9.2/go.mod h1:RwKglItbNcQq9UW6Vm2aOwOWXR6wTxUNjvuXcCIyJkU=
+github.com/loopholelabs/frisbee-go v0.9.3-0.20240908010142-de266dc4d891 h1:4Kt+YuiZYGB7BIZtS41Ynzjbqq7ekLLKB6yYhd1upcA=
+github.com/loopholelabs/frisbee-go v0.9.3-0.20240908010142-de266dc4d891/go.mod h1:RwKglItbNcQq9UW6Vm2aOwOWXR6wTxUNjvuXcCIyJkU=
 github.com/loopholelabs/logging v0.3.0 h1:Rfo9fGdBk4nwbNdiLrVUF7JkYoC67dvGDodkRe81xF0=
 github.com/loopholelabs/logging v0.3.0/go.mod h1:uRDUydiqPqKbZkb0WoQ3dfyAcJ2iOMhxdEafZssLVv0=
 github.com/loopholelabs/polyglot/v2 v2.0.2 h1:v308fg2ZKSvkKDnWgBnDvvmiu4YypCxcDe5Ih5GUVnY=

diff --git a/pkg/generator/generator.go b/pkg/generator/generator.go
@@ -114,7 +114,6 @@ func (g *Generator) Generate(req *pluginpb.CodeGeneratorRequest) (res *pluginpb.
 			"package":             packageName,
 			"requiredImports":     requiredImports,
 			"serviceImports":      serviceImports,
-			"methodImports":       methodImports,
 			"streamMethodImports": streamMethodImports,
 			"numServices":         numServices,
 			"numMethods":          numMethods,

diff --git a/pkg/generator/imports.go b/pkg/generator/imports.go
@@ -6,11 +6,12 @@ var (
 	requiredImports = []string{
 		"errors",
 		"net",
+		"sync",
+		"context",
 		"github.com/loopholelabs/polyglot/v2",
 	}
 
 	serviceImports = []string{
-		"context",
 		"crypto/tls",
 		"github.com/loopholelabs/frisbee-go",
 		"github.com/loopholelabs/frisbee-go/pkg/packet",
@@ -21,8 +22,4 @@ var (
 		"sync/atomic",
 		"io",
 	}
-
-	methodImports = []string{
-		"sync",
-	}
 )
diff --git a/pkg/generator/test/server.go b/pkg/generator/test/server.go
@@ -23,7 +23,7 @@ func (s svc) Echo(_ context.Context, request *Request) (*Response, error) {
 	}}, nil
 }
 
-func (s svc) EchoStream(srv *EchoStreamServer) error {
+func (s svc) EchoStream(_ context.Context, srv *EchoStreamServer) error {
 	for {
 		request, err := srv.Recv()
 		if err == io.EOF {
@@ -52,7 +52,7 @@ func (s svc) Testy(_ context.Context, _ *SearchResponse) (*StockPricesWrapper, e
 	panic("not implemented")
 }
 
-func (s svc) Search(req *SearchResponse, srv *SearchServer) error {
+func (s svc) Search(_ context.Context, req *SearchResponse, srv *SearchServer) error {
 	assert.Equal(s.t, 1, len(req.Results))
 	for i := 0; i < 10; i++ {
 		err := srv.Send(&Response{Message: "Hello World", Test: &Data{
@@ -64,7 +64,7 @@ func (s svc) Search(req *SearchResponse, srv *SearchServer) error {
 	return srv.CloseSend()
 }
 
-func (s svc) Upload(srv *UploadServer) error {
+func (s svc) Upload(_ context.Context, srv *UploadServer) error {
 	received := 0
 	for {
 		res, err := srv.Recv()

diff --git a/pkg/generator/test/test.frpc.go b/pkg/generator/test/test.frpc.go
diff --git a/templates/imports.templ b/templates/imports.templ
@@ -8,11 +8,6 @@ import (
     "{{$im}}"
 {{end -}}
 {{end}}
-{{ if .numMethods }}
-{{range $im := .methodImports -}}
-    "{{$im}}"
-{{end -}}
-{{end}}
 {{ if .numStreamMethods -}}
     {{range $im := .streamMethodImports -}}
         "{{$im}}"

diff --git a/templates/interfaces.templ b/templates/interfaces.templ
@@ -11,9 +11,9 @@ type {{ CamelCaseName .Name }} interface {
     {{ range $i, $v := MakeIterable .Methods.Len -}}
     {{ $method := $.Methods.Get $i -}}
     {{ if $method.IsStreamingClient }}
-    {{ CamelCaseName $method.Name }} (srv *{{ CamelCaseName $method.Name }}Server) error
+    {{ CamelCaseName $method.Name }} (context.Context, *{{ CamelCaseName $method.Name }}Server) error
     {{ else if $method.IsStreamingServer }}
-    {{ CamelCaseName $method.Name }} (req *{{ CamelCase $method.Input.FullName }}, srv *{{ CamelCaseName $method.Name }}Server) error
+    {{ CamelCaseName $method.Name }} (context.Context, *{{ CamelCase $method.Input.FullName }}, *{{ CamelCaseName $method.Name }}Server) error
     {{else -}}
     {{ CamelCaseName $method.Name }} (context.Context, *{{ CamelCase $method.Input.FullName }}) (*{{ CamelCase $method.Output.FullName }}, error)
     {{end -}}