🍦 serve: Rework NewServer and jig CLI for stacked dirs (#27)

Rework NewServer and jig CLI for stacked dirs. Instead of
specifying --protoset=FILE and --method-dir=DIR separately rework `jig
server` to load all *.pb files as FileDescriptorSets from the given
directories and search for all method stubs in the given directories too.

Allow for specifying several directories where the first found match is
served. To this end create StackedFS, an implementation of fs.FS, which
combines several fs.FS as if they were a single one, again returning
the first found match.

Pull-Request: #27

README.md

@@ -31,7 +31,7 @@ gRPC server receives a call to that method.
 
 You can generate skeleton jsonnet method definitions using `jig bones`:
 
-    jig bones --proto-set=service.pb --method-dir=dir
+    jig bones --proto-set=dir/service.pb --method-dir=dir
 
 Request protobuf messages are marshaled to JSON and passed to the jsonnet method
 definition function as the `input` parameter. If the method is a unary,
@@ -100,8 +100,7 @@ according to the [protojson] encoding rules.
 
 To serve these jsonnet methods, run:
 
-    jig serve --proto-set=service.pb --method-dir=dir
-
+    jig serve <dir>
 
 [gRPC status]: https://www.grpc.io/docs/guides/error/
 [protojson]: https://developers.google.com/protocol-buffers/docs/proto3#json
@@ -115,7 +114,7 @@ Build and start jig on the test data:
 
     . ./bin/activate-hermit
     make install
-    jig serve --proto-set pb/greet/greeter.pb --method-dir testdata
+    jig serve testdata pb/greet
 
 in a second terminal call it with:
 

main.go

@@ -17,10 +17,11 @@ type config struct {
 }
 
 type cmdServe struct {
-	LogLevel  serve.LogLevel `help:"Server logging level." default:"error"`
-	ProtoSet  string         `short:"p" help:"Protoset .pb file containing service and deps" required:""`
-	MethodDir string         `short:"m" default:"." help:"Directory containing method definitions"`
-	Listen    string         `short:"l" default:"localhost:8080" help:"TCP listen address"`
+	ProtoSet []string       `short:"p" help:"Protoset .pb files containing service and deps"`
+	LogLevel serve.LogLevel `help:"Server logging level." default:"error"`
+	Listen   string         `short:"l" default:"localhost:8080" help:"TCP listen address"`
+
+	Dirs []string `arg:"" help:"Directory containing method definitions and protoset .pb file"`
 }
 
 type cmdBones struct {
@@ -38,8 +39,9 @@ func main() {
 }
 
 func (cs *cmdServe) Run() error {
-	logger := serve.NewLogger(os.Stderr, cs.LogLevel)
-	s, err := serve.NewServer(cs.MethodDir, cs.ProtoSet, serve.WithLogger(logger))
+	withLogger := serve.WithLogger(serve.NewLogger(os.Stderr, cs.LogLevel))
+	withDirs := serve.WithDirs(cs.Dirs...)
+	s, err := serve.NewServer(withDirs, withLogger)
 	if err != nil {
 		return err
 	}

serve/server.go

@@ -8,6 +8,7 @@ import (
 	"io/fs"
 	"net"
 	"os"
+	"strings"
 
 	"foxygo.at/jig/reflection"
 	"github.com/google/go-jsonnet"
@@ -27,9 +28,26 @@ type Option func(s *Server) error
 // MakeVM is a constructor for a jsonnet VMs, exposed for custom configuration.
 type MakeVM func() *jsonnet.VM
 
-func WithFS(fs fs.FS) Option {
+func WithFS(dirs ...fs.FS) Option {
 	return func(s *Server) error {
-		s.fs = fs
+		s.fs = append(s.fs, dirs...)
+		return nil
+	}
+}
+
+func WithDirs(dirs ...string) Option {
+	return func(s *Server) error {
+		for _, dir := range dirs {
+			vfs := os.DirFS(dir)
+			s.fs = append(s.fs, vfs)
+		}
+		return nil
+	}
+}
+
+func WithProtosets(protosets ...string) Option {
+	return func(s *Server) error {
+		s.protosets = append(s.protosets, protosets...)
 		return nil
 	}
 }
@@ -49,31 +67,31 @@ func WithVM(makeVM MakeVM) Option {
 }
 
 type Server struct {
-	methodDir string
-	protoSet  string
-
-	log     Logger
-	methods map[string]method
-	gs      *grpc.Server
-	files   *protoregistry.Files
-	fs      fs.FS
-	makeVM  MakeVM
+	log       Logger
+	methods   map[string]method
+	gs        *grpc.Server
+	files     *protoregistry.Files
+	fs        stackedFS
+	protosets []string
+	makeVM    MakeVM
 }
 
 var errUnknownHandler = errors.New("Unknown handler")
 
 // NewServer creates a new Server. Its API is currently unstable.
-func NewServer(methodDir, protoSet string, options ...Option) (*Server, error) {
+func NewServer(options ...Option) (*Server, error) {
 	s := &Server{
-		methodDir: methodDir,
-		protoSet:  protoSet,
-		log:       NewLogger(os.Stderr, LogLevelError),
+		files: new(protoregistry.Files),
+		log:   NewLogger(os.Stderr, LogLevelError),
 	}
 	for _, opt := range options {
 		if err := opt(s); err != nil {
 			return nil, err
 		}
 	}
+	if len(s.fs) == 0 {
+		return nil, fmt.Errorf("missing directory")
+	}
 	if err := s.loadMethods(); err != nil {
 		return nil, err
 	}
@@ -104,38 +122,73 @@ func (s *Server) Stop() {
 }
 
 func (s *Server) loadMethods() error {
-	var b []byte
-	var err error
-	methodFS := s.fs
-	if s.fs != nil {
-		b, err = fs.ReadFile(s.fs, s.protoSet)
-	} else {
-		b, err = os.ReadFile(s.protoSet)
-		methodFS = os.DirFS(s.methodDir)
+	if err := s.loadProtosets(); err != nil {
+		return err
 	}
+
+	s.methods = make(map[string]method)
+	s.files.RangeFiles(func(fd protoreflect.FileDescriptor) bool {
+		sds := fd.Services()
+		for i := 0; i < sds.Len(); i++ {
+			mds := sds.Get(i).Methods()
+			for j := 0; j < mds.Len(); j++ {
+				m := newMethod(mds.Get(j), s.fs, s.makeVM)
+				s.methods[m.fullMethod()] = m
+			}
+		}
+		return true
+	})
+	return nil
+}
+
+func (s *Server) loadProtosets() error {
+	seen := map[string]bool{}
+	for _, protoset := range s.protosets {
+		b, err := os.ReadFile(protoset)
+		if err != nil {
+			return err
+		}
+		if err := s.addFiles(b, seen); err != nil {
+			return err
+		}
+	}
+
+	matches, err := fs.Glob(s.fs, "*.pb")
 	if err != nil {
 		return err
 	}
+	for _, match := range matches {
+		if strings.HasPrefix(match, "_") {
+			continue
+		}
+		b, err := fs.ReadFile(s.fs, match)
+		if err != nil {
+			return err
+		}
+		if err := s.addFiles(b, seen); err != nil {
+			return err
+		}
+	}
+	return nil
+}
 
+func (s *Server) addFiles(b []byte, seen map[string]bool) error {
 	fds := &descriptorpb.FileDescriptorSet{}
 	if err := proto.Unmarshal(b, fds); err != nil {
 		return err
 	}
-
-	s.files, err = protodesc.NewFiles(fds)
+	files, err := protodesc.NewFiles(fds)
 	if err != nil {
 		return err
 	}
-
-	s.methods = make(map[string]method)
-	s.files.RangeFiles(func(fd protoreflect.FileDescriptor) bool {
-		sds := fd.Services()
-		for i := 0; i < sds.Len(); i++ {
-			mds := sds.Get(i).Methods()
-			for j := 0; j < mds.Len(); j++ {
-				m := newMethod(mds.Get(j), methodFS, s.makeVM)
-				s.methods[m.fullMethod()] = m
-			}
+	files.RangeFiles(func(fd protoreflect.FileDescriptor) bool {
+		if seen[fd.Path()] {
+			return true
+		}
+		seen[fd.Path()] = true
+		err := s.files.RegisterFile(fd)
+		if err != nil {
+			s.log.Errorf("cannot register %q: %v", fd.FullName(), err)
 		}
 		return true
 	})
@@ -181,8 +234,8 @@ type TestServer struct {
 
 // NewTestServer starts and returns a new TestServer.
 // The caller should call Stop when finished, to shut it down.
-func NewTestServer(methodDir, protoSet string, options ...Option) *TestServer {
-	s, err := NewServer(methodDir, protoSet, options...)
+func NewTestServer(options ...Option) *TestServer {
+	s, err := NewServer(options...)
 	if err != nil {
 		panic(fmt.Sprintf("failed to create TestServer: %v", err))
 	}

serve/server_test.go

@@ -3,15 +3,19 @@ package serve
 import (
 	"bytes"
 	"embed"
+	"io"
 	"io/fs"
+	"os"
 	"testing"
 
 	"foxygo.at/jig/internal/client"
 	"github.com/stretchr/testify/require"
 )
 
 func newTestServer() *TestServer {
-	return NewTestServer("testdata/greet", "testdata/greet/greeter.pb")
+	lopOpt := WithLogger(NewLogger(io.Discard, LogLevelError))
+	fsOpt := WithFS(os.DirFS("testdata/greet"))
+	return NewTestServer(lopOpt, fsOpt)
 }
 
 type testCase struct {
@@ -123,7 +127,7 @@ var embedFS embed.FS
 func TestGreeterEmbedFS(t *testing.T) {
 	methodFS, err := fs.Sub(embedFS, "testdata/greet")
 	require.NoError(t, err)
-	ts := NewTestServer("NOT-RELEVANT-METHOD-DIR", "greeter.pb", WithFS(methodFS))
+	ts := NewTestServer(WithFS(methodFS))
 	defer ts.Stop()
 
 	c, err := client.New(ts.Addr())

serve/stackedfs.go

@@ -0,0 +1,39 @@
+package serve
+
+import (
+	"io/fs"
+	"sort"
+)
+
+type stackedFS []fs.FS
+
+// Open opens the the first occurrence of named file.
+func (s stackedFS) Open(name string) (f fs.File, err error) {
+	for _, vfs := range s {
+		if f, err = vfs.Open(name); err == nil {
+			return f, nil
+		}
+	}
+	return nil, err
+}
+
+// ReadDir combines all files on the stack, sorted by stack order first
+// and alphabetically within the stack second. Directories are not merged.
+func (s stackedFS) ReadDir(name string) (result []fs.DirEntry, err error) {
+	seen := map[string]bool{}
+	for _, vfs := range s {
+		entries, err := fs.ReadDir(vfs, name)
+		if err != nil {
+			return nil, err
+		}
+		byName := func(i, j int) bool { return entries[i].Name() < entries[j].Name() }
+		sort.Slice(entries, byName)
+		for _, entry := range entries {
+			if !seen[entry.Name()] {
+				seen[entry.Name()] = true
+				result = append(result, entry)
+			}
+		}
+	}
+	return result, nil
+}

serve/stackedfs_test.go

@@ -0,0 +1,35 @@
+package serve
+
+import (
+	"io/fs"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestStackedFS(t *testing.T) {
+	aFS := os.DirFS("testdata/stackedfs/a")
+	bFS := os.DirFS("testdata/stackedfs/b")
+	stacked := stackedFS{aFS, bFS}
+
+	b, err := fs.ReadFile(stacked, "1.txt")
+	require.NoError(t, err)
+	require.Equal(t, "1 in a\n", string(b))
+	b, err = fs.ReadFile(stacked, "3.txt")
+	require.NoError(t, err)
+	require.Equal(t, "3 in b\n", string(b))
+	b, err = fs.ReadFile(stacked, "2.txt")
+	require.NoError(t, err)
+	require.Equal(t, "2 in a\n", string(b))
+
+	entries, err := fs.ReadDir(stacked, ".")
+	require.NoError(t, err)
+	require.Equal(t, 4, len(entries))
+	var got []string
+	for _, e := range entries {
+		got = append(got, e.Name())
+	}
+	want := []string{"1.txt", "2.txt", "4.txt", "3.txt"}
+	require.Equal(t, want, got)
+}

serve/testdata/stackedfs/a/1.txt

@@ -0,0 +1 @@
+1 in a

serve/testdata/stackedfs/a/2.txt

@@ -0,0 +1 @@
+2 in a

serve/testdata/stackedfs/a/4.txt

@@ -0,0 +1 @@
+4 in a

serve/testdata/stackedfs/b/2.txt

@@ -0,0 +1 @@
+2 in b

serve/testdata/stackedfs/b/3.txt

@@ -0,0 +1 @@
+3 in b