Database Persistence for Cached Outputs

When multiple Flyte workers attempt to execute the same task simultaneously, the system must coordinate to prevent redundant computation and ensure that once a result is produced, it is reliably stored and retrievable. The Flyte Cache Service manages this through a dedicated persistence layer that tracks both the final outputs of tasks and the active reservations held by workers currently computing those outputs.

Repository Architecture

The persistence layer is structured around the Repository interface defined in cache_service/repository/interfaces/repository.go. This interface acts as a composite provider for two specialized repositories:

type Repository interface {
	CachedOutputRepo() CachedOutputRepo
	ReservationRepo() ReservationRepo
}

Flyte provides a SQL-based implementation of these interfaces using the sqlx library. The implementation is split into two primary classes: CachedOutputRepo and ReservationRepo, both located in the cache_service/repository/impl package. These repositories interact with a relational database (typically PostgreSQL) to maintain the state of the cache.

Cached Output Persistence

The CachedOutputRepo manages the cache_service_outputs table, which stores the mapping between a unique task execution key and the location of its results in object storage.

Data Model

The CachedOutput model in cache_service/repository/models/cached_output.go represents a completed task's result:

type CachedOutput struct {
	Key         string    `db:"key"`
	OutputURI   string    `db:"output_uri"`
	Metadata    []byte    `db:"metadata"`
	LastUpdated time.Time `db:"last_updated"`
	CreatedAt   time.Time `db:"created_at"`
	UpdatedAt   time.Time `db:"updated_at"`
}

Upsert Logic

To handle high-concurrency environments where multiple workers might finish the same task at nearly the same time, the Put method in cache_service/repository/impl/cached_output.go uses an INSERT ... ON CONFLICT (UPSERT) strategy. This ensures that the most recent output URI and metadata are preserved without causing primary key violations:

func (r *CachedOutputRepo) Put(ctx context.Context, output *models.CachedOutput) error {
	_, err := r.db.ExecContext(ctx,
		`INSERT INTO cache_service_outputs (key, output_uri, metadata, last_updated, created_at, updated_at)
		 VALUES ($1, $2, $3, $4, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
		 ON CONFLICT (key) DO UPDATE SET
		   output_uri = EXCLUDED.output_uri,
		   metadata = EXCLUDED.metadata,
		   last_updated = EXCLUDED.last_updated,
		   updated_at = CURRENT_TIMESTAMP`,
		output.Key, output.OutputURI, output.Metadata, output.LastUpdated)
	return err
}

Task Reservations and Concurrency

The ReservationRepo manages the cache_service_reservations table, which coordinates active work. When a worker starts a task that is not yet cached, it "reserves" the task key to signal to other workers that the work is already in progress.

Atomic Reservation Claims

The core of Flyte's coordination logic lies in the UpdateIfExpiredOrOwned method in cache_service/repository/impl/reservation.go. This method allows a worker to either extend its own existing reservation (heartbeat) or claim a reservation that has expired. The atomicity of the SQL UPDATE statement prevents race conditions where two workers might both believe they have claimed the same task:

func (r *ReservationRepo) UpdateIfExpiredOrOwned(ctx context.Context, reservation *models.Reservation, now time.Time) error {
	result, err := r.db.ExecContext(ctx,
		`UPDATE cache_service_reservations
		 SET owner_id = $1, heartbeat_seconds = $2, expires_at = $3, updated_at = $4
		 WHERE key = $5 AND (expires_at <= $6 OR owner_id = $7)`,
		reservation.OwnerID, reservation.HeartbeatSeconds, reservation.ExpiresAt, now,
		reservation.Key, now, reservation.OwnerID)
	if err != nil {
		return err
	}
	rows, _ := result.RowsAffected()
	if rows == 0 {
		return repositoryerrors.ErrReservationNotClaimable
	}
	return nil
}

If RowsAffected is zero, it indicates that another worker holds a valid, unexpired reservation, and the current worker should wait or back off.

Reservation Expiration

Reservations are governed by the Reservation model, which includes an ExpiresAt timestamp and a HeartbeatSeconds interval.

type Reservation struct {
	Key              string    `db:"key"`
	OwnerID          string    `db:"owner_id"`
	HeartbeatSeconds int64     `db:"heartbeat_seconds"`
	ExpiresAt        time.Time `db:"expires_at"`
	CreatedAt        time.Time `db:"created_at"`
	UpdatedAt        time.Time `db:"updated_at"`
}

The Manager (in cache_service/manager/manager.go) typically calculates the ExpiresAt value by multiplying the heartbeat interval by a grace period multiplier (defaulting to 3). This provides a buffer for network latency or transient worker delays before the reservation is considered abandoned and available for other workers to claim.

Integration with the Cache Manager

The Manager class acts as the primary consumer of these repositories. It orchestrates the high-level caching workflow:

1. Check Cache: Calls CachedOutputRepo.Get.
2. Coordinate Work: If the cache is empty, it calls ReservationRepo.Create or UpdateIfExpiredOrOwned to secure a lock.
3. Heartbeat: Periodically calls UpdateIfExpiredOrOwned to keep the reservation alive while the task runs.
4. Commit Result: Once the task finishes, it calls CachedOutputRepo.Put to store the result and ReservationRepo.DeleteByKeyAndOwner to release the lock.

This separation of concerns allows Flyte to maintain a consistent view of task status across a distributed set of workers using standard SQL ACID properties.