github.com/squaremo/docker@v1.3.2-0.20150516120342-42cfc9554972/pkg/truncindex/truncindex.go (about)

     1  package truncindex
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"strings"
     7  	"sync"
     8  
     9  	"github.com/tchap/go-patricia/patricia"
    10  )
    11  
    12  var (
    13  	ErrEmptyPrefix     = errors.New("Prefix can't be empty")
    14  	ErrAmbiguousPrefix = errors.New("Multiple IDs found with provided prefix")
    15  )
    16  
    17  // TruncIndex allows the retrieval of string identifiers by any of their unique prefixes.
    18  // This is used to retrieve image and container IDs by more convenient shorthand prefixes.
    19  type TruncIndex struct {
    20  	sync.RWMutex
    21  	trie *patricia.Trie
    22  	ids  map[string]struct{}
    23  }
    24  
    25  // NewTruncIndex creates a new TruncIndex and initializes with a list of IDs
    26  func NewTruncIndex(ids []string) (idx *TruncIndex) {
    27  	idx = &TruncIndex{
    28  		ids: make(map[string]struct{}),
    29  
    30  		// Change patricia max prefix per node length,
    31  		// because our len(ID) always 64
    32  		trie: patricia.NewTrie(patricia.MaxPrefixPerNode(64)),
    33  	}
    34  	for _, id := range ids {
    35  		idx.addID(id)
    36  	}
    37  	return
    38  }
    39  
    40  func (idx *TruncIndex) addID(id string) error {
    41  	if strings.Contains(id, " ") {
    42  		return fmt.Errorf("illegal character: ' '")
    43  	}
    44  	if id == "" {
    45  		return ErrEmptyPrefix
    46  	}
    47  	if _, exists := idx.ids[id]; exists {
    48  		return fmt.Errorf("id already exists: '%s'", id)
    49  	}
    50  	idx.ids[id] = struct{}{}
    51  	if inserted := idx.trie.Insert(patricia.Prefix(id), struct{}{}); !inserted {
    52  		return fmt.Errorf("failed to insert id: %s", id)
    53  	}
    54  	return nil
    55  }
    56  
    57  // Add adds a new ID to the TruncIndex
    58  func (idx *TruncIndex) Add(id string) error {
    59  	idx.Lock()
    60  	defer idx.Unlock()
    61  	if err := idx.addID(id); err != nil {
    62  		return err
    63  	}
    64  	return nil
    65  }
    66  
    67  // Delete removes an ID from the TruncIndex. If there are multiple IDs
    68  // with the given prefix, an error is thrown.
    69  func (idx *TruncIndex) Delete(id string) error {
    70  	idx.Lock()
    71  	defer idx.Unlock()
    72  	if _, exists := idx.ids[id]; !exists || id == "" {
    73  		return fmt.Errorf("no such id: '%s'", id)
    74  	}
    75  	delete(idx.ids, id)
    76  	if deleted := idx.trie.Delete(patricia.Prefix(id)); !deleted {
    77  		return fmt.Errorf("no such id: '%s'", id)
    78  	}
    79  	return nil
    80  }
    81  
    82  // Get retrieves an ID from the TruncIndex. If there are multiple IDs
    83  // with the given prefix, an error is thrown.
    84  func (idx *TruncIndex) Get(s string) (string, error) {
    85  	if s == "" {
    86  		return "", ErrEmptyPrefix
    87  	}
    88  	var (
    89  		id string
    90  	)
    91  	subTreeVisitFunc := func(prefix patricia.Prefix, item patricia.Item) error {
    92  		if id != "" {
    93  			// we haven't found the ID if there are two or more IDs
    94  			id = ""
    95  			return ErrAmbiguousPrefix
    96  		}
    97  		id = string(prefix)
    98  		return nil
    99  	}
   100  
   101  	idx.RLock()
   102  	defer idx.RUnlock()
   103  	if err := idx.trie.VisitSubtree(patricia.Prefix(s), subTreeVisitFunc); err != nil {
   104  		return "", err
   105  	}
   106  	if id != "" {
   107  		return id, nil
   108  	}
   109  	return "", fmt.Errorf("no such id: %s", s)
   110  }