From bb6c0c2de7f4738cc6329b9832c7a77429e359e6 Mon Sep 17 00:00:00 2001 From: Maksym Pavlenko Date: Sun, 11 Apr 2021 19:32:46 -0700 Subject: [PATCH] Add more bolt utils This PR extracts out a few more bolt utils from the codebase for easier management. Signed-off-by: Maksym Pavlenko --- metadata/boltutil/helpers.go | 88 ++++++++++++++++++++++++++++++++++++ metadata/containers.go | 69 +++++----------------------- 2 files changed, 99 insertions(+), 58 deletions(-) diff --git a/metadata/boltutil/helpers.go b/metadata/boltutil/helpers.go index 94af315f9..6995917c9 100644 --- a/metadata/boltutil/helpers.go +++ b/metadata/boltutil/helpers.go @@ -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 +} diff --git a/metadata/containers.go b/metadata/containers.go index 09b0d203d..29d3e22c1 100644 --- a/metadata/containers.go +++ b/metadata/containers.go @@ -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,15 +369,8 @@ func writeContainer(bkt *bolt.Bucket, container *containers.Container) error { return err } - if container.Spec != nil { - spec, err := container.Spec.Marshal() - if err != nil { - return err - } - - if err := bkt.Put(bucketKeySpec, spec); err != nil { - return err - } + if err := boltutil.WriteAny(bkt, bucketKeySpec, container.Spec); err != nil { + return err } for _, v := range [][2][]byte{ @@ -424,33 +398,12 @@ func writeContainer(bkt *bolt.Bucket, container *containers.Container) error { return err } - if len(container.Extensions) > 0 { - ebkt, err := bkt.CreateBucketIfNotExists(bucketKeyExtensions) - if err != nil { - return err - } - - for name, ext := range container.Extensions { - p, err := proto.Marshal(&ext) - if err != nil { - return err - } - - if err := ebkt.Put([]byte(name), p); err != nil { - return err - } - } + if err := boltutil.WriteExtensions(bkt, container.Extensions); 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 - } + if err := boltutil.WriteAny(rbkt, bucketKeyOptions, container.Runtime.Options); err != nil { + return err } return boltutil.WriteLabels(bkt, container.Labels)