From 0a3d5e8b4014d266f130bd52ec844ea295e8b919 Mon Sep 17 00:00:00 2001
From: Simone Gotti <simone.gotti@gmail.com>
Date: Wed, 4 Mar 2015 12:10:42 +0100
Subject: [PATCH] cas: use the db in ResolveKey.

As the db is the primary and coherent information source, use it instead of
diskv in ResolveKey.

This will also help in future diskv locking as it removes the need for complex
locking inside ResolveKey.
---
 cas/aciinfo.go  | 20 ++++++++++++++++++++
 cas/cas.go      | 21 +++++++++++----------
 cas/cas_test.go | 13 ++++++++++++-
 3 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/cas/aciinfo.go b/cas/aciinfo.go
index bfad1fe..cfea3db 100644
--- a/cas/aciinfo.go
+++ b/cas/aciinfo.go
@@ -27,6 +27,26 @@ func NewACIInfo(blobKey string, latest bool, t time.Time) *ACIInfo {
 	}
 }
 
+// GetAciInfosWithKeyPrefix returns all the ACIInfos with a blobkey starting with the given prefix.
+func GetACIInfosWithKeyPrefix(tx *sql.Tx, prefix string) ([]*ACIInfo, error) {
+	aciinfos := []*ACIInfo{}
+	rows, err := tx.Query("SELECT * from aciinfo WHERE hasPrefix(blobkey, $1)", prefix)
+	if err != nil {
+		return nil, err
+	}
+	for rows.Next() {
+		aciinfo := &ACIInfo{}
+		if err := rows.Scan(&aciinfo.BlobKey, &aciinfo.AppName, &aciinfo.ImportTime, &aciinfo.Latest); err != nil {
+			return nil, err
+		}
+		aciinfos = append(aciinfos, aciinfo)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return aciinfos, err
+}
+
 // GetAciInfosWithAppName returns all the ACIInfos for a given appname. found will be
 // false if no aciinfo exists.
 func GetACIInfosWithAppName(tx *sql.Tx, appname string) ([]*ACIInfo, bool, error) {
diff --git a/cas/cas.go b/cas/cas.go
index 5f72660..712db5c 100644
--- a/cas/cas.go
+++ b/cas/cas.go
@@ -148,23 +148,24 @@ func (ds Store) ResolveKey(key string) (string, error) {
 		key = key[:lenKey]
 	}
 
-	cancel := make(chan struct{})
-	var k string
-	keyCount := 0
-	for k = range ds.stores[blobType].KeysPrefix(key, cancel) {
-		keyCount++
-		if keyCount > 1 {
-			close(cancel)
-			break
-		}
+	aciInfos := []*ACIInfo{}
+	err := ds.db.Do(func(tx *sql.Tx) error {
+		var err error
+		aciInfos, err = GetACIInfosWithKeyPrefix(tx, key)
+		return err
+	})
+	if err != nil {
+		return "", fmt.Errorf("error retrieving ACI Infos: %v", err)
 	}
+
+	keyCount := len(aciInfos)
 	if keyCount == 0 {
 		return "", fmt.Errorf("no keys found")
 	}
 	if keyCount != 1 {
 		return "", fmt.Errorf("ambiguous key: %q", key)
 	}
-	return k, nil
+	return aciInfos[0].BlobKey, nil
 }
 
 func (ds Store) ReadStream(key string) (io.ReadCloser, error) {
diff --git a/cas/cas_test.go b/cas/cas_test.go
index 4bf2571..af5b309 100644
--- a/cas/cas_test.go
+++ b/cas/cas_test.go
@@ -17,12 +17,14 @@ package cas
 import (
 	"archive/tar"
 	"bytes"
+	"database/sql"
 	"encoding/hex"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/coreos/rocket/Godeps/_workspace/src/github.com/appc/spec/schema/types"
 	"github.com/coreos/rocket/pkg/aci"
@@ -160,7 +162,16 @@ func TestResolveKey(t *testing.T) {
 		str2key("67147019a5b56f5e2ee01e989a8aa4787f56b8445960be2d8678391cf111009bc0780f31001fd181a2b61507547aee4caa44cda4b8bdb238d0e4ba830069ed2c"),
 	}
 	for _, d := range data {
-		if err := ds.WriteStream(d.String(), d); err != nil {
+		// Save aciinfo
+		err := ds.db.Do(func(tx *sql.Tx) error {
+			aciinfo := &ACIInfo{
+				BlobKey:    d.String(),
+				AppName:    "example.com/app",
+				ImportTime: time.Now(),
+			}
+			return WriteACIInfo(tx, aciinfo)
+		})
+		if err != nil {
 			t.Fatalf("error writing to store: %v", err)
 		}
 	}