From 2ab70f21ac6d9d03716f2dd87577d6b3b0bc6d92 Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcgstyle.net>
Date: Fri, 18 Aug 2017 15:16:33 -0700
Subject: [PATCH] Add content and snapshot store references

Update database object to hold reference to the
data stores.

Signed-off-by: Derek McGowan <derek@mcgstyle.net>
---
 metadata/content.go       | 13 ++++++++++---
 metadata/content_test.go  |  2 +-
 metadata/db.go            | 11 ++++++++++-
 metadata/snapshot.go      | 17 +++++++++++++----
 metadata/snapshot_test.go |  2 +-
 5 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/metadata/content.go b/metadata/content.go
index e93e09292..7e4572666 100644
--- a/metadata/content.go
+++ b/metadata/content.go
@@ -25,10 +25,17 @@ type contentStore struct {
 // NewContentStore returns a namespaced content store using an existing
 // content store interface.
 func NewContentStore(db *DB, cs content.Store) content.Store {
-	return &contentStore{
-		Store: cs,
-		db:    db,
+	db.storeL.Lock()
+	defer db.storeL.Unlock()
+
+	if db.cs == nil {
+		db.cs = &contentStore{
+			Store: cs,
+			db:    db,
+		}
 	}
+	return db.cs
+
 }
 
 func (cs *contentStore) Info(ctx context.Context, dgst digest.Digest) (content.Info, error) {
diff --git a/metadata/content_test.go b/metadata/content_test.go
index 8d23c33f9..fbddff2fc 100644
--- a/metadata/content_test.go
+++ b/metadata/content_test.go
@@ -23,7 +23,7 @@ func createContentStore(ctx context.Context, root string) (content.Store, func()
 		return nil, nil, err
 	}
 
-	return NewContentStore(db, cs), func() error {
+	return NewContentStore(NewDB(db), cs), func() error {
 		return db.Close()
 	}, nil
 }
diff --git a/metadata/db.go b/metadata/db.go
index a6009d83a..958983f4e 100644
--- a/metadata/db.go
+++ b/metadata/db.go
@@ -1,14 +1,23 @@
 package metadata
 
-import "github.com/boltdb/bolt"
+import (
+	"sync"
+
+	"github.com/boltdb/bolt"
+)
 
 type DB struct {
 	db *bolt.DB
+
+	storeL sync.Mutex
+	ss     map[string]*snapshotter
+	cs     *contentStore
 }
 
 func NewDB(db *bolt.DB) *DB {
 	return &DB{
 		db: db,
+		ss: map[string]*snapshotter{},
 	}
 }
 
diff --git a/metadata/snapshot.go b/metadata/snapshot.go
index d83c4676d..aeae2519a 100644
--- a/metadata/snapshot.go
+++ b/metadata/snapshot.go
@@ -25,11 +25,20 @@ type snapshotter struct {
 // NewSnapshotter returns a new Snapshotter which namespaces the given snapshot
 // using the provided name and metadata store.
 func NewSnapshotter(db *DB, name string, sn snapshot.Snapshotter) snapshot.Snapshotter {
-	return &snapshotter{
-		Snapshotter: sn,
-		name:        name,
-		db:          db,
+	db.storeL.Lock()
+	defer db.storeL.Unlock()
+
+	ss, ok := db.ss[name]
+	if !ok {
+		ss = &snapshotter{
+			Snapshotter: sn,
+			name:        name,
+			db:          db,
+		}
+		db.ss[name] = ss
 	}
+
+	return ss
 }
 
 func createKey(id uint64, namespace, key string) string {
diff --git a/metadata/snapshot_test.go b/metadata/snapshot_test.go
index 37800dc40..e7ad54e3e 100644
--- a/metadata/snapshot_test.go
+++ b/metadata/snapshot_test.go
@@ -28,7 +28,7 @@ func newSnapshotter(ctx context.Context, root string) (snapshot.Snapshotter, fun
 		return nil, nil, err
 	}
 
-	sn := NewSnapshotter(db, "naive", snapshotter)
+	sn := NewSnapshotter(NewDB(db), "naive", snapshotter)
 
 	return sn, func() error {
 		return db.Close()