feat: persistent task worker queues

We create a new boltdb for manager to persistent task worker queues in boltdb.
Manager will load the task queue from the database and sync new tasks with database.

Signed-off-by: Yadong Ding <ding_yadong@foxmail.com>

pkg/content/local.go

@@ -27,6 +27,7 @@ import (
 	"github.com/containerd/containerd/remotes"
 	"github.com/goharbor/acceleration-service/pkg/config"
 	"github.com/goharbor/acceleration-service/pkg/remote"
+	"github.com/goharbor/acceleration-service/pkg/task"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 )
@@ -46,6 +47,9 @@ func NewLocalProvider(cfg *config.Config, platformMC platforms.MatchComparer) (P
 	if err := os.MkdirAll(contentDir, 0755); err != nil {
 		return nil, nil, errors.Wrap(err, "create local provider work directory")
 	}
+	if err := task.Manager.Init(cfg.Provider.WorkDir); err != nil {
+		return nil, nil, errors.Wrap(err, "task manager init")
+	}
 	content, err := NewContent(contentDir, cfg.Provider.WorkDir, cfg.Provider.GCPolicy.Threshold,
 		cfg.EnableRemoteCache(), cfg.Provider.CacheSize, cfg.Host)
 

pkg/task/manager.go

@@ -15,30 +15,37 @@
 package task
 
 import (
+	"encoding/json"
+	"path/filepath"
 	"sort"
 	"sync"
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/pkg/errors"
+	bolt "go.etcd.io/bbolt"
 )
 
+var bucketObjectTasks = []byte("tasks")
+
 const taskMaximumKeepPeriod = time.Hour * 24
 
 const StatusProcessing = "PROCESSING"
 const StatusCompleted = "COMPLETED"
 const StatusFailed = "FAILED"
 
 type Task struct {
-	ID       string     `json:"id"`
-	Created  time.Time  `json:"created"`
-	Finished *time.Time `json:"finished"`
-	Source   string     `json:"source"`
-	Status   string     `json:"status"`
-	Reason   string     `json:"reason"`
+	ID       string    `json:"id"`
+	Created  time.Time `json:"created"`
+	Finished time.Time `json:"finished"`
+	Source   string    `json:"source"`
+	Status   string    `json:"status"`
+	Reason   string    `json:"reason"`
 }
 
 type manager struct {
 	mutex sync.Mutex
+	db    *bolt.DB
 	tasks map[string]*Task
 }
 
@@ -59,24 +66,89 @@ func (t *Task) IsExpired() bool {
 	return false
 }
 
+// Init manager supported by boltdb.
+func (m *manager) Init(workDir string) error {
+	bdb, err := bolt.Open(filepath.Join(workDir, "task.db"), 0655, nil)
+	if err != nil {
+		return errors.Wrap(err, "create task database")
+	}
+	m.db = bdb
+	return m.initDatabase()
+}
+
+// initDatabase loads tasks from the database into memory.
+func (m *manager) initDatabase() error {
+	return m.db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket([]byte("tasks"))
+		if bucket == nil {
+			return nil
+		}
+
+		return bucket.ForEach(func(k, v []byte) error {
+			var task Task
+			if err := json.Unmarshal(v, &task); err != nil {
+				return err
+			}
+			m.tasks[task.ID] = &task
+			return nil
+		})
+	})
+}
+
+// updateBucket updates task in bucket and creates a new bucket if it doesn't already exist.
+func (m *manager) updateBucket(task *Task) error {
+	return m.db.Update(func(tx *bolt.Tx) error {
+		bucket, err := tx.CreateBucketIfNotExists(bucketObjectTasks)
+		if err != nil {
+			return err
+		}
+
+		taskJSON, err := json.Marshal(task)
+		if err != nil {
+			return err
+		}
+
+		if err := bucket.Put([]byte(task.ID), taskJSON); err != nil {
+			return err
+		}
+
+		return nil
+	})
+}
+
+// deleteBucket deletes a task in bucket
+func (m *manager) deleteBucket(taskId string) error {
+	return m.db.Update(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(bucketObjectTasks)
+		if bucket == nil {
+			return nil
+		}
+		return bucket.Delete([]byte(taskId))
+	})
+}
+
+// Create new task
 func (m *manager) Create(source string) string {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
 	id := uuid.NewString()
 
-	m.tasks[id] = &Task{
-		ID:       id,
-		Created:  time.Now(),
-		Finished: nil,
-		Source:   source,
-		Status:   StatusProcessing,
-		Reason:   "",
+	task := &Task{
+		ID:      id,
+		Created: time.Now(),
+		Source:  source,
+		Status:  StatusProcessing,
+		Reason:  "",
 	}
 
+	m.updateBucket(task)
+	m.tasks[id] = task
+
 	return id
 }
 
+// Finish a task
 func (m *manager) Finish(id string, err error) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
@@ -90,14 +162,15 @@ func (m *manager) Finish(id string, err error) {
 		} else {
 			task.Status = StatusCompleted
 		}
-		now := time.Now()
-		task.Finished = &now
+		task.Finished = time.Now()
 	}
+	m.updateBucket(task)
 
 	// Evict expired tasks.
 	for id, task := range m.tasks {
 		if task.IsExpired() {
 			delete(m.tasks, id)
+			m.deleteBucket(id)
 		}
 	}
 }