Add bounded scale-to-zero leases

server/cmd/api/api/scaletozero.go

@@ -2,11 +2,17 @@ package api
 
 import (
 	"context"
+	"errors"
+	"regexp"
+	"time"
 
 	"github.com/kernel/kernel-images/server/lib/logger"
 	oapi "github.com/kernel/kernel-images/server/lib/oapi"
+	"github.com/kernel/kernel-images/server/lib/scaletozero"
 )
 
+var scaleToZeroLeaseIDPattern = regexp.MustCompile(`^[A-Za-z0-9_-]{1,128}$`)
+
 func (s *ApiService) DisableScaleToZero(ctx context.Context, _ oapi.DisableScaleToZeroRequestObject) (oapi.DisableScaleToZeroResponseObject, error) {
 	if err := s.stz.Pin(ctx); err != nil {
 		logger.FromContext(ctx).Error("failed to disable scale-to-zero", "err", err)
@@ -22,3 +28,29 @@ func (s *ApiService) EnableScaleToZero(ctx context.Context, _ oapi.EnableScaleTo
 	}
 	return oapi.EnableScaleToZero204Response{}, nil
 }
+
+func (s *ApiService) AcquireScaleToZeroLease(ctx context.Context, req oapi.AcquireScaleToZeroLeaseRequestObject) (oapi.AcquireScaleToZeroLeaseResponseObject, error) {
+	if !scaleToZeroLeaseIDPattern.MatchString(req.LeaseId) || req.Params.TtlSeconds < 1 || req.Params.TtlSeconds > 300 {
+		return oapi.AcquireScaleToZeroLease400JSONResponse{BadRequestErrorJSONResponse: oapi.BadRequestErrorJSONResponse{Message: "invalid scale-to-zero lease"}}, nil
+	}
+	ttl := time.Duration(req.Params.TtlSeconds) * time.Second
+	if err := s.stz.AcquireLease(ctx, req.LeaseId, ttl); err != nil {
+		if errors.Is(err, scaletozero.ErrLeaseLimit) {
+			return oapi.AcquireScaleToZeroLease409JSONResponse{ConflictErrorJSONResponse: oapi.ConflictErrorJSONResponse{Message: scaletozero.ErrLeaseLimit.Error()}}, nil
+		}
+		logger.FromContext(ctx).Error("failed to acquire scale-to-zero lease", "err", err, "lease_id", req.LeaseId)
+		return oapi.AcquireScaleToZeroLease500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to acquire scale-to-zero lease"}}, nil
+	}
+	return oapi.AcquireScaleToZeroLease204Response{}, nil
+}
+
+func (s *ApiService) ReleaseScaleToZeroLease(ctx context.Context, req oapi.ReleaseScaleToZeroLeaseRequestObject) (oapi.ReleaseScaleToZeroLeaseResponseObject, error) {
+	if !scaleToZeroLeaseIDPattern.MatchString(req.LeaseId) {
+		return oapi.ReleaseScaleToZeroLease400JSONResponse{BadRequestErrorJSONResponse: oapi.BadRequestErrorJSONResponse{Message: "invalid scale-to-zero lease"}}, nil
+	}
+	if err := s.stz.ReleaseLease(ctx, req.LeaseId); err != nil {
+		logger.FromContext(ctx).Error("failed to release scale-to-zero lease", "err", err, "lease_id", req.LeaseId)
+		return oapi.ReleaseScaleToZeroLease500JSONResponse{InternalErrorJSONResponse: oapi.InternalErrorJSONResponse{Message: "failed to release scale-to-zero lease"}}, nil
+	}
+	return oapi.ReleaseScaleToZeroLease204Response{}, nil
+}

server/cmd/api/api/scaletozero_test.go

@@ -0,0 +1,63 @@
+package api
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	oapi "github.com/kernel/kernel-images/server/lib/oapi"
+	"github.com/kernel/kernel-images/server/lib/scaletozero"
+	"github.com/stretchr/testify/require"
+)
+
+func TestScaleToZeroLeaseHandlers(t *testing.T) {
+	t.Parallel()
+
+	svc := &ApiService{stz: scaletozero.NewNoopController()}
+	acquired, err := svc.AcquireScaleToZeroLease(t.Context(), oapi.AcquireScaleToZeroLeaseRequestObject{
+		LeaseId: "metro-lease",
+		Params:  oapi.AcquireScaleToZeroLeaseParams{TtlSeconds: 30},
+	})
+	require.NoError(t, err)
+	require.IsType(t, oapi.AcquireScaleToZeroLease204Response{}, acquired)
+
+	released, err := svc.ReleaseScaleToZeroLease(t.Context(), oapi.ReleaseScaleToZeroLeaseRequestObject{LeaseId: "metro-lease"})
+	require.NoError(t, err)
+	require.IsType(t, oapi.ReleaseScaleToZeroLease204Response{}, released)
+}
+
+func TestScaleToZeroLeaseHandlerReturnsConflictAtLimit(t *testing.T) {
+	t.Parallel()
+
+	svc := &ApiService{stz: &leaseLimitController{NoopController: scaletozero.NewNoopController()}}
+	response, err := svc.AcquireScaleToZeroLease(t.Context(), oapi.AcquireScaleToZeroLeaseRequestObject{
+		LeaseId: "metro-lease",
+		Params:  oapi.AcquireScaleToZeroLeaseParams{TtlSeconds: 30},
+	})
+	require.NoError(t, err)
+	require.IsType(t, oapi.AcquireScaleToZeroLease409JSONResponse{}, response)
+}
+
+func TestScaleToZeroLeaseHandlersRejectInvalidInput(t *testing.T) {
+	t.Parallel()
+
+	svc := &ApiService{stz: scaletozero.NewNoopController()}
+	acquired, err := svc.AcquireScaleToZeroLease(t.Context(), oapi.AcquireScaleToZeroLeaseRequestObject{
+		LeaseId: "invalid/lease",
+		Params:  oapi.AcquireScaleToZeroLeaseParams{TtlSeconds: 301},
+	})
+	require.NoError(t, err)
+	require.IsType(t, oapi.AcquireScaleToZeroLease400JSONResponse{}, acquired)
+
+	released, err := svc.ReleaseScaleToZeroLease(t.Context(), oapi.ReleaseScaleToZeroLeaseRequestObject{LeaseId: "invalid/lease"})
+	require.NoError(t, err)
+	require.IsType(t, oapi.ReleaseScaleToZeroLease400JSONResponse{}, released)
+}
+
+type leaseLimitController struct {
+	*scaletozero.NoopController
+}
+
+func (*leaseLimitController) AcquireLease(context.Context, string, time.Duration) error {
+	return scaletozero.ErrLeaseLimit
+}