Add more bolt utils

This PR extracts out a few more bolt utils from the codebase for easier management.

Signed-off-by: Maksym Pavlenko <pavlenko.maksym@gmail.com>
This commit is contained in:
Maksym Pavlenko 2021-04-11 19:32:46 -07:00
parent 2486eb8ff8
commit bb6c0c2de7
2 changed files with 99 additions and 58 deletions

View File

@ -19,6 +19,8 @@ package boltutil
import (
"time"
"github.com/gogo/protobuf/proto"
"github.com/gogo/protobuf/types"
"github.com/pkg/errors"
bolt "go.etcd.io/bbolt"
)
@ -28,6 +30,7 @@ var (
bucketKeyLabels = []byte("labels")
bucketKeyCreatedAt = []byte("createdat")
bucketKeyUpdatedAt = []byte("updatedat")
bucketKeyExtensions = []byte("extensions")
)
// ReadLabels reads the labels key from the bucket
@ -145,3 +148,88 @@ func WriteTimestamps(bkt *bolt.Bucket, created, updated time.Time) error {
return nil
}
// WriteExtensions will write a KV map to the given bucket,
// where `K` is a string key and `V` is a protobuf's Any type that represents a generic extension.
func WriteExtensions(bkt *bolt.Bucket, extensions map[string]types.Any) error {
if len(extensions) == 0 {
return nil
}
ebkt, err := bkt.CreateBucketIfNotExists(bucketKeyExtensions)
if err != nil {
return err
}
for name, ext := range extensions {
p, err := proto.Marshal(&ext)
if err != nil {
return err
}
if err := ebkt.Put([]byte(name), p); err != nil {
return err
}
}
return nil
}
// ReadExtensions will read back a map of extensions from the given bucket, previously written by WriteExtensions
func ReadExtensions(bkt *bolt.Bucket) (map[string]types.Any, error) {
var (
extensions = make(map[string]types.Any)
ebkt = bkt.Bucket(bucketKeyExtensions)
)
if ebkt == nil {
return extensions, nil
}
if err := ebkt.ForEach(func(k, v []byte) error {
var t types.Any
if err := proto.Unmarshal(v, &t); err != nil {
return err
}
extensions[string(k)] = t
return nil
}); err != nil {
return nil, err
}
return extensions, nil
}
// WriteAny write a protobuf's Any type to the bucket
func WriteAny(bkt *bolt.Bucket, name []byte, any *types.Any) error {
if any == nil {
return nil
}
data, err := proto.Marshal(any)
if err != nil {
return err
}
if err := bkt.Put(name, data); err != nil {
return err
}
return nil
}
// ReadAny reads back protobuf's Any type from the bucket
func ReadAny(bkt *bolt.Bucket, name []byte) (*types.Any, error) {
bytes := bkt.Get(name)
if bytes == nil {
return nil, nil
}
out := types.Any{}
if err := proto.Unmarshal(bytes, &out); err != nil {
return nil, errors.Wrap(err, "failed to unmarshal any")
}
return &out, nil
}

View File

@ -336,16 +336,11 @@ func readContainer(container *containers.Container, bkt *bolt.Bucket) error {
container.Runtime.Name = string(n)
}
obkt := rbkt.Get(bucketKeyOptions)
if obkt == nil {
return nil
}
var any types.Any
if err := proto.Unmarshal(obkt, &any); err != nil {
any, err := boltutil.ReadAny(rbkt, bucketKeyOptions)
if err != nil {
return err
}
container.Runtime.Options = &any
container.Runtime.Options = any
case string(bucketKeySpec):
var any types.Any
if err := proto.Unmarshal(v, &any); err != nil {
@ -357,22 +352,8 @@ func readContainer(container *containers.Container, bkt *bolt.Bucket) error {
case string(bucketKeySnapshotter):
container.Snapshotter = string(v)
case string(bucketKeyExtensions):
ebkt := bkt.Bucket(bucketKeyExtensions)
if ebkt == nil {
return nil
}
extensions := make(map[string]types.Any)
if err := ebkt.ForEach(func(k, v []byte) error {
var a types.Any
if err := proto.Unmarshal(v, &a); err != nil {
return err
}
extensions[string(k)] = a
return nil
}); err != nil {
extensions, err := boltutil.ReadExtensions(bkt)
if err != nil {
return err
}
@ -388,17 +369,10 @@ func writeContainer(bkt *bolt.Bucket, container *containers.Container) error {
return err
}
if container.Spec != nil {
spec, err := container.Spec.Marshal()
if err != nil {
if err := boltutil.WriteAny(bkt, bucketKeySpec, container.Spec); err != nil {
return err
}
if err := bkt.Put(bucketKeySpec, spec); err != nil {
return err
}
}
for _, v := range [][2][]byte{
{bucketKeyImage, []byte(container.Image)},
{bucketKeySnapshotter, []byte(container.Snapshotter)},
@ -424,34 +398,13 @@ func writeContainer(bkt *bolt.Bucket, container *containers.Container) error {
return err
}
if len(container.Extensions) > 0 {
ebkt, err := bkt.CreateBucketIfNotExists(bucketKeyExtensions)
if err != nil {
if err := boltutil.WriteExtensions(bkt, container.Extensions); err != nil {
return err
}
for name, ext := range container.Extensions {
p, err := proto.Marshal(&ext)
if err != nil {
if err := boltutil.WriteAny(rbkt, bucketKeyOptions, container.Runtime.Options); err != nil {
return err
}
if err := ebkt.Put([]byte(name), p); err != nil {
return err
}
}
}
if container.Runtime.Options != nil {
data, err := proto.Marshal(container.Runtime.Options)
if err != nil {
return err
}
if err := rbkt.Put(bucketKeyOptions, data); err != nil {
return err
}
}
return boltutil.WriteLabels(bkt, container.Labels)
}