feat: add batch services support (#91)

Co-authored-by: Ryan Cartwright <[email protected]>

Author: jyecusch

Makefile

@@ -48,6 +48,7 @@ generate:
 	go run github.com/golang/mock/mockgen github.com/nitrictech/nitric/core/pkg/proto/storage/v1 StorageClient > mocks/storage.go
 	go run github.com/golang/mock/mockgen github.com/nitrictech/nitric/core/pkg/proto/secrets/v1 SecretManagerClient > mocks/secrets.go
 	go run github.com/golang/mock/mockgen github.com/nitrictech/nitric/core/pkg/proto/topics/v1 TopicsClient > mocks/topics.go
+	go run github.com/golang/mock/mockgen github.com/nitrictech/nitric/core/pkg/proto/batch/v1 BatchClient > mocks/batch.go
 	go run github.com/golang/mock/mockgen -package mock_v1 google.golang.org/grpc ClientConnInterface > mocks/grpc_clientconn.go
 
 # Runs tests for coverage upload to codecov.io

go.mod

@@ -10,7 +10,7 @@ require (
 	github.com/golangci/golangci-lint v1.61.0
 	github.com/google/addlicense v1.1.1
 	github.com/missionMeteora/toolkit v0.0.0-20170713173850-88364e3ef8cc
-	github.com/nitrictech/nitric/core v0.0.0-20240915234849-42c1e482ddab
+	github.com/nitrictech/nitric/core v0.0.0-20241003062412-76ea6275fb0b
 	github.com/nitrictech/protoutils v0.0.0-20220321044654-02667a814cdf
 	github.com/onsi/ginkgo v1.16.5
 	github.com/onsi/gomega v1.34.2
@@ -118,7 +118,7 @@ require (
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/polyfloyd/go-errorlint v1.6.0 // indirect
 	github.com/prometheus/client_golang v1.14.0 // indirect
-	github.com/prometheus/client_model v0.5.0 // indirect
+	github.com/prometheus/client_model v0.6.0 // indirect
 	github.com/prometheus/common v0.37.0 // indirect
 	github.com/prometheus/procfs v0.8.0 // indirect
 	github.com/quasilyte/go-ruleguard v0.4.3-0.20240823090925-0fe6f58b47b1 // indirect
@@ -220,5 +220,5 @@ require (
 	go-simpler.org/musttag v0.12.2 // indirect
 	go-simpler.org/sloglint v0.7.2 // indirect
 	go.uber.org/automaxprocs v1.5.3 // indirect
-	go.uber.org/goleak v1.2.1 // indirect
+	go.uber.org/goleak v1.3.0 // indirect
 )

go.sum

@@ -424,8 +424,10 @@ github.com/nishanths/exhaustive v0.12.0 h1:vIY9sALmw6T/yxiASewa4TQcFsVYZQQRUQJhK
 github.com/nishanths/exhaustive v0.12.0/go.mod h1:mEZ95wPIZW+x8kC4TgC+9YCUgiST7ecevsVDTgc2obs=
 github.com/nishanths/predeclared v0.2.2 h1:V2EPdZPliZymNAn79T8RkNApBjMmVKh5XRpLm/w98Vk=
 github.com/nishanths/predeclared v0.2.2/go.mod h1:RROzoN6TnGQupbC+lqggsOlcgysk3LMK/HI84Mp280c=
-github.com/nitrictech/nitric/core v0.0.0-20240915234849-42c1e482ddab h1:59YTNUa6en385Y21SwMTaixbe31SwgzjkjEhxx0RuwQ=
-github.com/nitrictech/nitric/core v0.0.0-20240915234849-42c1e482ddab/go.mod h1:N274XVBjYhGEQoT42baWM6/lETBQYQhqPpqUuk2gmLc=
+github.com/nitrictech/nitric/core v0.0.0-20240913000004-5d21c28b00ba h1:ZIPl9waqhbqw3xB2+zpUI2T1kEHyMkOnZZMt6tdrEUM=
+github.com/nitrictech/nitric/core v0.0.0-20240913000004-5d21c28b00ba/go.mod h1:4LQH9hea9rX+0A+8G47NRk5nZuXCDqiwci1aZsHAkU8=
+github.com/nitrictech/nitric/core v0.0.0-20241003062412-76ea6275fb0b h1:ImQFk66gRM3v9A6qmPImOiV3HJMDAX93X5rplMKn6ok=
+github.com/nitrictech/nitric/core v0.0.0-20241003062412-76ea6275fb0b/go.mod h1:9bQnYPqLzq8CcPk5MHT3phg19CWJhDlFOfdIv27lwwM=
 github.com/nitrictech/protoutils v0.0.0-20220321044654-02667a814cdf h1:8MB8W8ylM8sCM2COGfiO39/tB6BTdiawLszaUGCNL5w=
 github.com/nitrictech/protoutils v0.0.0-20220321044654-02667a814cdf/go.mod h1:b2lzk2a4o1bvSrSCE6yvTldHuXCJymuDVhdMJGOSslw=
 github.com/nunnatsa/ginkgolinter v0.16.2 h1:8iLqHIZvN4fTLDC0Ke9tbSZVcyVHoBs0HIbnVSxfHJk=
@@ -481,6 +483,7 @@ github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:
 github.com/prometheus/client_model v0.2.0/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
 github.com/prometheus/client_model v0.5.0 h1:VQw1hfvPvk3Uv6Qf29VrPF32JB6rtbgI6cYPYQjL0Qw=
 github.com/prometheus/client_model v0.5.0/go.mod h1:dTiFglRmd66nLR9Pv9f0mZi7B7fk5Pm3gvsjB5tr+kI=
+github.com/prometheus/client_model v0.6.0/go.mod h1:NTQHnmxFpouOD0DpvP4XujX3CdOAGQPoaGhyTchlyt8=
 github.com/prometheus/common v0.4.1/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y86RQel1bk4=
 github.com/prometheus/common v0.10.0/go.mod h1:Tlit/dnDKsSWFlCLTWaA1cyBgKHSMdTB80sz/V91rCo=
 github.com/prometheus/common v0.26.0/go.mod h1:M7rCNAaPfAosfx8veZJCuw84e35h3Cfd9VFqTh1DIvc=
@@ -642,6 +645,7 @@ go.uber.org/automaxprocs v1.5.3 h1:kWazyxZUrS3Gs4qUpbwo5kEIMGe/DAvi5Z4tl2NW4j8=
 go.uber.org/automaxprocs v1.5.3/go.mod h1:eRbA25aqJrxAbsLO0xy5jVwPt7FQnRgjW+efnwa1WM0=
 go.uber.org/goleak v1.2.1 h1:NBol2c7O1ZokfZ0LEU9K6Whx/KnwvepVetCUhtKja4A=
 go.uber.org/goleak v1.2.1/go.mod h1:qlT2yGI9QafXHhZZLxlSuNsMw3FFLxBr+tBRlmO1xH4=
+go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
 go.uber.org/multierr v1.8.0 h1:dg6GjLku4EH+249NNmoIciG9N/jURbDG+pFlTkhzIC8=
 go.uber.org/multierr v1.8.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak=
 go.uber.org/zap v1.24.0 h1:FiJd5l1UOLj0wCgbSE0rwwXHzEdAZS6hiiSnxJN/D60=

mocks/batch.go

@@ -47,8 +47,10 @@ type Route interface {
 	ApiName() string
 }
 
-type Handler = handlers.Handler[Ctx]
-type Middleware = handlers.Middleware[Ctx]
+type (
+	Handler    = handlers.Handler[Ctx]
+	Middleware = handlers.Middleware[Ctx]
+)
 
 type route struct {
 	path       string

nitric/apis/api.go

@@ -0,0 +1,144 @@
+// Copyright 2021 Nitric Technologies Pty Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package batch
+
+import (
+	"fmt"
+
+	"github.com/nitrictech/go-sdk/internal/handlers"
+	"github.com/nitrictech/go-sdk/nitric/workers"
+	batchpb "github.com/nitrictech/nitric/core/pkg/proto/batch/v1"
+	v1 "github.com/nitrictech/nitric/core/pkg/proto/resources/v1"
+)
+
+// JobPermission defines the available permissions on a job
+type JobPermission string
+
+type Handler = handlers.Handler[Ctx]
+
+const (
+	// JobSubmit is required to call Submit on a job.
+	JobSubmit JobPermission = "submit"
+)
+
+type JobReference interface {
+	// Allow requests the given permissions to the job.
+	Allow(permission JobPermission, permissions ...JobPermission) *BatchClient
+
+	// Handler will register and start the job task handler that will be called for all task submitted to this job.
+	// Valid function signatures for middleware are:
+	//
+	//	func()
+	//	func() error
+	//	func(*batch.Ctx)
+	//	func(*batch.Ctx) error
+	//	Handler[batch.Ctx]
+	Handler(handler interface{}, options ...HandlerOption)
+}
+
+type jobReference struct {
+	name         string
+	manager      *workers.Manager
+	registerChan <-chan workers.RegisterResult
+}
+
+// NewJob creates a new job resource with the give name.
+func NewJob(name string) JobReference {
+	job := &jobReference{
+		name:    name,
+		manager: workers.GetDefaultManager(),
+	}
+
+	job.registerChan = job.manager.RegisterResource(&v1.ResourceDeclareRequest{
+		Id: &v1.ResourceIdentifier{
+			Type: v1.ResourceType_Job,
+			Name: name,
+		},
+		Config: &v1.ResourceDeclareRequest_Job{
+			Job: &v1.JobResource{},
+		},
+	})
+
+	return job
+}
+
+func (j *jobReference) Allow(permission JobPermission, permissions ...JobPermission) *BatchClient {
+	allPerms := append([]JobPermission{permission}, permissions...)
+
+	actions := []v1.Action{}
+	for _, perm := range allPerms {
+		switch perm {
+		case JobSubmit:
+			actions = append(actions, v1.Action_JobSubmit)
+		default:
+			panic(fmt.Errorf("JobPermission %s unknown", perm))
+		}
+	}
+
+	registerResult := <-j.registerChan
+	if registerResult.Err != nil {
+		panic(registerResult.Err)
+	}
+
+	err := j.manager.RegisterPolicy(registerResult.Identifier, actions...)
+	if err != nil {
+		panic(err)
+	}
+
+	client, err := NewBatchClient(j.name)
+	if err != nil {
+		panic(err)
+	}
+
+	return client
+}
+
+func (j *jobReference) Handler(handler interface{}, opts ...HandlerOption) {
+	options := &handlerOptions{}
+
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	registrationRequest := &batchpb.RegistrationRequest{
+		JobName:      j.name,
+		Requirements: &batchpb.JobResourceRequirements{},
+	}
+
+	if options.cpus != nil {
+		registrationRequest.Requirements.Cpus = *options.cpus
+	}
+
+	if options.memory != nil {
+		registrationRequest.Requirements.Memory = *options.memory
+	}
+
+	if options.gpus != nil {
+		registrationRequest.Requirements.Gpus = *options.gpus
+	}
+
+	typedHandler, err := handlers.HandlerFromInterface[Ctx](handler)
+	if err != nil {
+		panic(err)
+	}
+
+	jobOpts := &jobWorkerOpts{
+		RegistrationRequest: registrationRequest,
+		Handler:             typedHandler,
+	}
+
+	worker := newJobWorker(jobOpts)
+	j.manager.AddWorker("JobWorker:"+j.name, worker)
+}

nitric/batch/batch.go

@@ -0,0 +1,27 @@
+// Copyright 2023 Nitric Technologies Pty Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package batch_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+func TestBatch(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Batch (Jobs) Suite")
+}