Cache Serailize API

datacatalog_config.yaml

@@ -11,6 +11,8 @@ datacatalog:
   storage-prefix: "metadata"
   metrics-scope: "datacatalog"
   profiler-port: 10254
+  heartbeat-grace-period-multiplier: 3
+  max-reservation-heartbeat: 10s
 storage:
   connection:
     access-key: minio

go.mod

@@ -4,7 +4,7 @@ go 1.13
 
 require (
 	github.com/Selvatico/go-mocket v1.0.7
-	github.com/flyteorg/flyteidl v0.18.17
+	github.com/flyteorg/flyteidl v0.18.38
 	github.com/flyteorg/flytestdlib v0.3.13
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
 	github.com/golang/protobuf v1.4.3
@@ -18,3 +18,5 @@ require (
 	gorm.io/driver/postgres v1.1.0
 	gorm.io/gorm v1.21.9
 )
+
+replace github.com/flyteorg/flyteidl => github.com/hamersaw/flyteidl v0.19.26-0.20211103115633-100abab11c51

go.sum

@@ -175,8 +175,8 @@ github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5Kwzbycv
 github.com/fatih/color v1.10.0 h1:s36xzo75JdqLaaWoiEHk767eHiwo0598uUxyfiPkDsg=
 github.com/fatih/color v1.10.0/go.mod h1:ELkj/draVOlAH/xkhN6mQ50Qd0MPOk5AAr3maGEBuJM=
 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
-github.com/flyteorg/flyteidl v0.18.17 h1:74pPZ9PzITuzq+CgjMPb9EcFI5bVkf8mM5m4xmmlTmY=
-github.com/flyteorg/flyteidl v0.18.17/go.mod h1:b5Fq4Z8a5b0mF6pEwTd48ufvikUGVkWSjZiMT0ZtqKI=
+github.com/flyteorg/flyteidl v0.18.38 h1:XgAw9d2Q/UjWQyXbnZz/j4N6OVGDxr7jceden6PdCgY=
+github.com/flyteorg/flyteidl v0.18.38/go.mod h1:b5Fq4Z8a5b0mF6pEwTd48ufvikUGVkWSjZiMT0ZtqKI=
 github.com/flyteorg/flytestdlib v0.3.13 h1:5ioA/q3ixlyqkFh5kDaHgmPyTP/AHtqq1K/TIbVLUzM=
 github.com/flyteorg/flytestdlib v0.3.13/go.mod h1:Tz8JCECAbX6VWGwFT6cmEQ+RJpZ/6L9pswu3fzWs220=
 github.com/form3tech-oss/jwt-go v3.2.2+incompatible h1:TcekIExNqud5crz4xD2pavyTgWiPvpYe4Xau31I0PRk=
@@ -308,6 +308,9 @@ github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgf
 github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
 github.com/grpc-ecosystem/grpc-gateway v1.9.5/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
 github.com/grpc-ecosystem/grpc-gateway v1.12.2/go.mod h1:8XEsbTttt/W+VvjtQhLACqCisSPWTxCZ7sBRjU6iH9c=
+github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
+github.com/hamersaw/flyteidl v0.19.26-0.20211103115633-100abab11c51 h1:e9zvtfNKr+K84a7du4wJgC+MXYcEsA13yYGMsEGsjQs=
+github.com/hamersaw/flyteidl v0.19.26-0.20211103115633-100abab11c51/go.mod h1:576W2ViEyjTpT+kEVHAGbrTP3HARNUZ/eCwrNPmdx9U=
 github.com/hashicorp/consul/api v1.1.0/go.mod h1:VmuI/Lkw1nC05EYQWNKwWGbkg+FbDBtguAZLlVdkD9Q=
 github.com/hashicorp/consul/api v1.3.0/go.mod h1:MmDNSzIMUjNpY/mQ398R4bk2FnqQLoPndWW5VkKPlCE=
 github.com/hashicorp/consul/sdk v0.1.1/go.mod h1:VKf9jXwCTEY1QZP2MOLRhb5i/I/ssyNV1vwHyQBF0x8=
@@ -523,6 +526,7 @@ github.com/performancecopilot/speed v3.0.0+incompatible/go.mod h1:/CLtqpZ5gBg1M9
 github.com/peterbourgon/diskv v2.0.1+incompatible/go.mod h1:uqqh8zWWbv1HBMNONnaR/tNboyR3/BZd58JJSHlUSCU=
 github.com/pierrec/lz4 v1.0.2-0.20190131084431-473cd7ce01a1/go.mod h1:3/3N9NVKO0jef7pBehbT1qWhCMrIgbYNnFAZCqQ5LRc=
 github.com/pierrec/lz4 v2.0.5+incompatible/go.mod h1:pdkljMzZIN41W+lC3N2tnIh5sFi+IEE17M5jbnwPHcY=
+github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4/go.mod h1:N6UoU20jOqggOuDwUaBQpluzLNDqif3kq9z2wpdYEfQ=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
@@ -989,6 +993,7 @@ google.golang.org/genproto v0.0.0-20200312145019-da6875a35672/go.mod h1:55QSHmfG
 google.golang.org/genproto v0.0.0-20200331122359-1ee6d9798940/go.mod h1:55QSHmfGQM9UVYDPBsyGGes0y52j32PQ3BqQfXhyH3c=
 google.golang.org/genproto v0.0.0-20200430143042-b979b6f78d84/go.mod h1:55QSHmfGQM9UVYDPBsyGGes0y52j32PQ3BqQfXhyH3c=
 google.golang.org/genproto v0.0.0-20200511104702-f5ebc3bea380/go.mod h1:55QSHmfGQM9UVYDPBsyGGes0y52j32PQ3BqQfXhyH3c=
+google.golang.org/genproto v0.0.0-20200513103714-09dca8ec2884/go.mod h1:55QSHmfGQM9UVYDPBsyGGes0y52j32PQ3BqQfXhyH3c=
 google.golang.org/genproto v0.0.0-20200515170657-fc4c6c6a6587/go.mod h1:YsZOwe1myG/8QRHRsmBRE1LrgQY60beZKjly0O1fX9U=
 google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo=
 google.golang.org/genproto v0.0.0-20200618031413-b414f8b61790/go.mod h1:jDfRM7FcilCzHH/e9qn6dsT145K34l5v+OpcnNgKAAA=
@@ -1026,6 +1031,7 @@ google.golang.org/grpc v1.30.0/go.mod h1:N36X2cJ7JwdamYAgDz+s+rVMFjt3numwzf/HckM
 google.golang.org/grpc v1.31.0/go.mod h1:N36X2cJ7JwdamYAgDz+s+rVMFjt3numwzf/HckM8pak=
 google.golang.org/grpc v1.31.1/go.mod h1:N36X2cJ7JwdamYAgDz+s+rVMFjt3numwzf/HckM8pak=
 google.golang.org/grpc v1.32.0/go.mod h1:N36X2cJ7JwdamYAgDz+s+rVMFjt3numwzf/HckM8pak=
+google.golang.org/grpc v1.33.1/go.mod h1:fr5YgcSWrqhRRxogOsw7RzIpsmvOZ6IcH4kBYTpR3n0=
 google.golang.org/grpc v1.33.2/go.mod h1:JMHMWHQWaTccqQQlmk3MJZS+GWXOdAesneDmEnv2fbc=
 google.golang.org/grpc v1.34.0/go.mod h1:WotjhfgOW/POjDeRt8vscBtXq+2VjORFy659qA51WJ8=
 google.golang.org/grpc v1.35.0/go.mod h1:qjiiYl8FncCW8feJPdyg3v6XW24KsRHe+dy9BAGRRjU=

pkg/manager/impl/artifact_manager_test.go

@@ -90,8 +90,9 @@ func getTestArtifact() *datacatalog.Artifact {
 
 func newMockDataCatalogRepo() *mocks.DataCatalogRepo {
 	return &mocks.DataCatalogRepo{
-		MockDatasetRepo:  &mocks.DatasetRepo{},
-		MockArtifactRepo: &mocks.ArtifactRepo{},
+		MockDatasetRepo:     &mocks.DatasetRepo{},
+		MockArtifactRepo:    &mocks.ArtifactRepo{},
+		MockReservationRepo: &mocks.ReservationRepo{},
 	}
 }
 

pkg/manager/impl/reservation_manager.go

@@ -0,0 +1,212 @@
+package impl
+
+import (
+	"context"
+	"time"
+
+	"github.com/flyteorg/flytestdlib/logger"
+	"github.com/flyteorg/flytestdlib/promutils"
+	"github.com/flyteorg/flytestdlib/promutils/labeled"
+
+	"github.com/flyteorg/datacatalog/pkg/errors"
+	"github.com/flyteorg/datacatalog/pkg/repositories"
+	repo_errors "github.com/flyteorg/datacatalog/pkg/repositories/errors"
+	"github.com/flyteorg/datacatalog/pkg/repositories/models"
+	"github.com/flyteorg/datacatalog/pkg/repositories/transformers"
+
+	"github.com/flyteorg/datacatalog/pkg/manager/interfaces"
+
+	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog"
+)
+
+type reservationMetrics struct {
+	scope                        promutils.Scope
+	reservationAcquired          labeled.Counter
+	reservationReleased          labeled.Counter
+	reservationAlreadyInProgress labeled.Counter
+	acquireReservationFailure    labeled.Counter
+	releaseReservationFailure    labeled.Counter
+	reservationDoesNotExist      labeled.Counter
+}
+
+type NowFunc func() time.Time
+
+type reservationManager struct {
+	repo                           repositories.RepositoryInterface
+	heartbeatGracePeriodMultiplier time.Duration
+	maxHeartbeatInterval           time.Duration
+	now                            NowFunc
+	systemMetrics                  reservationMetrics
+}
+
+// Creates a new reservation manager with the specified properties
+func NewReservationManager(
+	repo repositories.RepositoryInterface,
+	heartbeatGracePeriodMultiplier time.Duration,
+	maxHeartbeatInterval time.Duration,
+	nowFunc NowFunc, // Easier to mock time.Time for testing
+	reservationScope promutils.Scope,
+) interfaces.ReservationManager {
+	systemMetrics := reservationMetrics{
+		scope: reservationScope,
+		reservationAcquired: labeled.NewCounter(
+			"reservation_acquired",
+			"Number of times a reservation was acquired",
+			reservationScope),
+		reservationReleased: labeled.NewCounter(
+			"reservation_released",
+			"Number of times a reservation was released",
+			reservationScope),
+		reservationAlreadyInProgress: labeled.NewCounter(
+			"reservation_already_in_progress",
+			"Number of times we try of acquire a reservation but the reservation is in progress",
+			reservationScope,
+		),
+		acquireReservationFailure: labeled.NewCounter(
+			"acquire_reservation_failure",
+			"Number of times we failed to acquire reservation",
+			reservationScope,
+		),
+		releaseReservationFailure: labeled.NewCounter(
+			"release_reservation_failure",
+			"Number of times we failed to release a reservation",
+			reservationScope,
+		),
+		reservationDoesNotExist: labeled.NewCounter(
+			"reservation_does_not_exist",
+			"Number of times we attempt to modify a reservation that does not exist",
+			reservationScope,
+		),
+	}
+
+	return &reservationManager{
+		repo:                           repo,
+		heartbeatGracePeriodMultiplier: heartbeatGracePeriodMultiplier,
+		maxHeartbeatInterval:           maxHeartbeatInterval,
+		now:                            nowFunc,
+		systemMetrics:                  systemMetrics,
+	}
+}
+
+// Attempt to acquire a reservation for the specified artifact. If there is not active reservation, successfully
+// acquire it. If you are the owner of the active reservation, extend it. If another owner, return the existing reservation.
+func (r *reservationManager) GetOrExtendReservation(ctx context.Context, request *datacatalog.GetOrExtendReservationRequest) (*datacatalog.GetOrExtendReservationResponse, error) {
+	reservationID := request.ReservationId
+
+	// Use minimum of maxHeartbeatInterval and requested heartbeat interval
+	heartbeatInterval := r.maxHeartbeatInterval
+	requestHeartbeatInterval := request.GetHeartbeatInterval()
+	if requestHeartbeatInterval != nil && requestHeartbeatInterval.AsDuration() < heartbeatInterval {
+		heartbeatInterval = requestHeartbeatInterval.AsDuration()
+	}
+
+	reservation, err := r.tryAcquireReservation(ctx, reservationID, request.OwnerId, heartbeatInterval)
+	if err != nil {
+		r.systemMetrics.acquireReservationFailure.Inc(ctx)
+		return nil, err
+	}
+
+	return &datacatalog.GetOrExtendReservationResponse{
+		Reservation: &reservation,
+	}, nil
+}
+
+// tryAcquireReservation will fetch the reservation first and only create/update
+// the reservation if it does not exist or has expired.
+// This is an optimization to reduce the number of writes to db. We always need
+// to do a GET here because we want to know who owns the reservation
+// and show it to users on the UI. However, the reservation is held by a single
+// task most of the times and there is no need to do a write.
+func (r *reservationManager) tryAcquireReservation(ctx context.Context, reservationID *datacatalog.ReservationID, ownerID string, heartbeatInterval time.Duration) (datacatalog.Reservation, error) {
+	repo := r.repo.ReservationRepo()
+	reservationKey := transformers.FromReservationID(reservationID)
+	repoReservation, err := repo.Get(ctx, reservationKey)
+
+	reservationExists := true
+	if err != nil {
+		if errors.IsDoesNotExistError(err) {
+			// Reservation does not exist yet so let's create one
+			reservationExists = false
+		} else {
+			return datacatalog.Reservation{}, err
+		}
+	}
+
+	now := r.now()
+	newRepoReservation := models.Reservation{
+		ReservationKey: reservationKey,
+		OwnerID:        ownerID,
+		ExpiresAt:      now.Add(heartbeatInterval * r.heartbeatGracePeriodMultiplier),
+	}
+
+	// Conditional upsert on reservation. Race conditions are handled
+	// within the reservation repository Create and Update function calls.
+	var repoErr error
+	if !reservationExists {
+		repoErr = repo.Create(ctx, newRepoReservation, now)
+	} else if repoReservation.ExpiresAt.Before(now) || repoReservation.OwnerID == ownerID {
+		repoErr = repo.Update(ctx, newRepoReservation, now)
+	} else {
+		logger.Debugf(ctx, "Reservation: %+v is held by %s", reservationKey, repoReservation.OwnerID)
+
+		reservation, err := transformers.CreateReservation(&repoReservation, heartbeatInterval)
+		if err != nil {
+			return reservation, err
+		}
+
+		r.systemMetrics.reservationAlreadyInProgress.Inc(ctx)
+		return reservation, nil
+	}
+
+	if repoErr != nil {
+		if repoErr.Error() == repo_errors.AlreadyExists {
+			// Looks like someone else tried to obtain the reservation
+			// at the same time and they won. Let's find out who won.
+			rsv1, err := repo.Get(ctx, reservationKey)
+			if err != nil {
+				return datacatalog.Reservation{}, err
+			}
+
+			reservation, err := transformers.CreateReservation(&rsv1, heartbeatInterval)
+			if err != nil {
+				return reservation, err
+			}
+
+			r.systemMetrics.reservationAlreadyInProgress.Inc(ctx)
+			return reservation, nil
+		}
+
+		return datacatalog.Reservation{}, repoErr
+	}
+
+	// Reservation has been acquired or extended without error
+	reservation, err := transformers.CreateReservation(&newRepoReservation, heartbeatInterval)
+	if err != nil {
+		return reservation, err
+	}
+
+	r.systemMetrics.reservationAlreadyInProgress.Inc(ctx)
+	return reservation, nil
+}
+
+// Release an active reservation with the specified owner. If one does not exist, gracefully return.
+func (r *reservationManager) ReleaseReservation(ctx context.Context, request *datacatalog.ReleaseReservationRequest) (*datacatalog.ReleaseReservationResponse, error) {
+	repo := r.repo.ReservationRepo()
+	reservationKey := transformers.FromReservationID(request.ReservationId)
+
+	err := repo.Delete(ctx, reservationKey, request.OwnerId)
+	if err != nil {
+		if errors.IsDoesNotExistError(err) {
+			logger.Warnf(ctx, "Reservation does not exist id: %+v, err %v", request.ReservationId, err)
+			r.systemMetrics.reservationDoesNotExist.Inc(ctx)
+			return &datacatalog.ReleaseReservationResponse{}, nil
+		}
+
+		logger.Errorf(ctx, "Failed to release reservation: %+v, err: %v", reservationKey, err)
+		r.systemMetrics.releaseReservationFailure.Inc(ctx)
+		return nil, err
+	}
+
+	r.systemMetrics.reservationReleased.Inc(ctx)
+	return &datacatalog.ReleaseReservationResponse{}, nil
+}