storj/pkg/piecestore/rpc/server/psdb/psdb.go

267 lines
5.2 KiB
Go

// Copyright (C) 2018 Storj Labs, Inc.
// See LICENSE for copying information.
package psdb
import (
"context"
"database/sql"
"flag"
"fmt"
"log"
"os"
"path/filepath"
"sync"
"time"
"go.uber.org/zap"
monkit "gopkg.in/spacemonkeygo/monkit.v2"
_ "github.com/mattn/go-sqlite3" // sqlite is weird and needs underscore
"storj.io/storj/pkg/piecestore"
pb "storj.io/storj/protos/piecestore"
)
var (
mon = monkit.Package()
defaultCheckInterval = flag.Duration("piecestore.ttl.check_interval",
time.Hour, "number of seconds to sleep between ttl checks")
)
// PSDB -- Piecestore database
type PSDB struct {
DB *sql.DB
mtx sync.Mutex
dataPath string
checkInterval time.Duration
}
// OpenPSDB -- opens PSDB at DBPath
func OpenPSDB(ctx context.Context, DataPath, DBPath string) (psdb *PSDB, err error) {
defer mon.Task()(&ctx)(&err)
if err = os.MkdirAll(filepath.Dir(DBPath), 0700); err != nil {
return nil, err
}
db, err := sql.Open("sqlite3", fmt.Sprintf("file:%s?cache=shared&mode=rwc&mutex=full", DBPath))
if err != nil {
return nil, err
}
defer func() {
if err != nil {
db.Close()
}
}()
tx, err := db.Begin()
if err != nil {
return nil, err
}
defer rollback(tx)
_, err = tx.Exec("CREATE TABLE IF NOT EXISTS `ttl` (`id` TEXT UNIQUE, `created` INT(10), `expires` INT(10));")
if err != nil {
return nil, err
}
_, err = tx.Exec("CREATE TABLE IF NOT EXISTS `bandwidth_agreements` (`agreement` BLOB, `signature` BLOB);")
if err != nil {
return nil, err
}
err = tx.Commit()
if err != nil {
return nil, err
}
return &PSDB{
DB: db,
dataPath: DataPath,
checkInterval: *defaultCheckInterval}, nil
}
// deleteEntries -- checks for and deletes expired TTL entries
func deleteEntriesFromFS(dir string, rows *sql.Rows) error {
for rows.Next() {
var expID string
err := rows.Scan(&expID)
if err != nil {
return err
}
// delete file on local machine
err = pstore.Delete(expID, dir)
if err != nil {
return err
}
log.Printf("Deleted file: %s\n", expID)
}
return rows.Err()
}
// Close the database
func (psdb *PSDB) Close() error {
return psdb.DB.Close()
}
// DeleteExpired checks for expired TTLs in the DB and removes data from both the DB and the FS
func (psdb *PSDB) DeleteExpired(ctx context.Context) (err error) {
defer mon.Task()(&ctx)(&err)
now := time.Now().Unix()
rows, err := psdb.DB.Query(fmt.Sprintf("SELECT id FROM ttl WHERE expires < %d AND expires > 0", now))
if err != nil {
return err
}
defer func() {
if rowErr := rows.Close(); rowErr != nil {
log.Printf("failed to close Rows: %s\n", rowErr)
}
}()
if err = deleteEntriesFromFS(psdb.dataPath, rows); err != nil {
return err
}
_, err = psdb.DB.Exec(fmt.Sprintf("DELETE FROM ttl WHERE expires < %d AND expires > 0", now))
if err != nil {
return err
}
return nil
}
// DeleteExpiredLoop will periodically run DeleteExpired
func (psdb *PSDB) DeleteExpiredLoop(ctx context.Context) (err error) {
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
time.Sleep(psdb.checkInterval)
err = ctx.Err()
if err != nil {
return err
}
err = psdb.DeleteExpired(ctx)
if err != nil {
zap.S().Errorf("failed checking entries: %+v", err)
}
}
}
}
// WriteBandwidthAllocToDB -- Insert bandwidth agreement into DB
func (psdb *PSDB) WriteBandwidthAllocToDB(ba *pb.RenterBandwidthAllocation) error {
psdb.mtx.Lock()
defer psdb.mtx.Unlock()
data := ba.GetData()
if data == nil {
return nil
}
tx, err := psdb.DB.Begin()
if err != nil {
return err
}
defer rollback(tx)
stmt, err := psdb.DB.Prepare(`INSERT INTO bandwidth_agreements (agreement, signature) VALUES (?, ?)`)
if err != nil {
return err
}
defer stmt.Close()
_, err = tx.Stmt(stmt).Exec(data, ba.GetSignature())
if err != nil {
return err
}
return tx.Commit()
}
// AddTTLToDB -- Insert TTL into database by id
func (psdb *PSDB) AddTTLToDB(id string, expiration int64) error {
psdb.mtx.Lock()
defer psdb.mtx.Unlock()
tx, err := psdb.DB.Begin()
if err != nil {
return err
}
defer rollback(tx)
stmt, err := psdb.DB.Prepare("INSERT or REPLACE INTO ttl (id, created, expires) VALUES (?, ?, ?)")
if err != nil {
return err
}
defer stmt.Close()
_, err = tx.Stmt(stmt).Exec(id, time.Now().Unix(), expiration)
if err != nil {
return err
}
return tx.Commit()
}
// GetTTLByID -- Find the TTL in the database by id and return it
func (psdb *PSDB) GetTTLByID(id string) (expiration int64, err error) {
psdb.mtx.Lock()
defer psdb.mtx.Unlock()
err = psdb.DB.QueryRow(fmt.Sprintf(`SELECT expires FROM ttl WHERE id="%s"`, id)).Scan(&expiration)
if err != nil {
return 0, err
}
return expiration, nil
}
// DeleteTTLByID -- Find the TTL in the database by id and delete it
func (psdb *PSDB) DeleteTTLByID(id string) error {
psdb.mtx.Lock()
defer psdb.mtx.Unlock()
tx, err := psdb.DB.Begin()
if err != nil {
return err
}
defer rollback(tx)
stmt, err := tx.Prepare("DELETE FROM ttl WHERE id=?")
if err != nil {
return err
}
defer stmt.Close()
_, err = tx.Stmt(stmt).Exec(id)
if err != nil {
return err
}
return tx.Commit()
}
func rollback(tx *sql.Tx) {
err := tx.Rollback()
if err != nil && err != sql.ErrTxDone {
log.Printf("%s\n", err)
}
}