github.com/demonoid81/moby@v0.0.0-20200517203328-62dd8e17c460/pkg/truncindex/truncindex.go (about)

     1  // Package truncindex provides a general 'index tree', used by Docker
     2  // in order to be able to reference containers by only a few unambiguous
     3  // characters of their id.
     4  package truncindex // import "github.com/demonoid81/moby/pkg/truncindex"
     5  
     6  import (
     7  	"errors"
     8  	"fmt"
     9  	"strings"
    10  	"sync"
    11  
    12  	"github.com/tchap/go-patricia/patricia"
    13  )
    14  
    15  var (
    16  	// ErrEmptyPrefix is an error returned if the prefix was empty.
    17  	ErrEmptyPrefix = errors.New("Prefix can't be empty")
    18  
    19  	// ErrIllegalChar is returned when a space is in the ID
    20  	ErrIllegalChar = errors.New("illegal character: ' '")
    21  
    22  	// ErrNotExist is returned when ID or its prefix not found in index.
    23  	ErrNotExist = errors.New("ID does not exist")
    24  )
    25  
    26  // ErrAmbiguousPrefix is returned if the prefix was ambiguous
    27  // (multiple ids for the prefix).
    28  type ErrAmbiguousPrefix struct {
    29  	prefix string
    30  }
    31  
    32  func (e ErrAmbiguousPrefix) Error() string {
    33  	return fmt.Sprintf("Multiple IDs found with provided prefix: %s", e.prefix)
    34  }
    35  
    36  // TruncIndex allows the retrieval of string identifiers by any of their unique prefixes.
    37  // This is used to retrieve image and container IDs by more convenient shorthand prefixes.
    38  type TruncIndex struct {
    39  	sync.RWMutex
    40  	trie *patricia.Trie
    41  	ids  map[string]struct{}
    42  }
    43  
    44  // NewTruncIndex creates a new TruncIndex and initializes with a list of IDs.
    45  func NewTruncIndex(ids []string) (idx *TruncIndex) {
    46  	idx = &TruncIndex{
    47  		ids: make(map[string]struct{}),
    48  
    49  		// Change patricia max prefix per node length,
    50  		// because our len(ID) always 64
    51  		trie: patricia.NewTrie(patricia.MaxPrefixPerNode(64)),
    52  	}
    53  	for _, id := range ids {
    54  		idx.addID(id)
    55  	}
    56  	return
    57  }
    58  
    59  func (idx *TruncIndex) addID(id string) error {
    60  	if strings.Contains(id, " ") {
    61  		return ErrIllegalChar
    62  	}
    63  	if id == "" {
    64  		return ErrEmptyPrefix
    65  	}
    66  	if _, exists := idx.ids[id]; exists {
    67  		return fmt.Errorf("id already exists: '%s'", id)
    68  	}
    69  	idx.ids[id] = struct{}{}
    70  	if inserted := idx.trie.Insert(patricia.Prefix(id), struct{}{}); !inserted {
    71  		return fmt.Errorf("failed to insert id: %s", id)
    72  	}
    73  	return nil
    74  }
    75  
    76  // Add adds a new ID to the TruncIndex.
    77  func (idx *TruncIndex) Add(id string) error {
    78  	idx.Lock()
    79  	defer idx.Unlock()
    80  	return idx.addID(id)
    81  }
    82  
    83  // Delete removes an ID from the TruncIndex. If there are multiple IDs
    84  // with the given prefix, an error is thrown.
    85  func (idx *TruncIndex) Delete(id string) error {
    86  	idx.Lock()
    87  	defer idx.Unlock()
    88  	if _, exists := idx.ids[id]; !exists || id == "" {
    89  		return fmt.Errorf("no such id: '%s'", id)
    90  	}
    91  	delete(idx.ids, id)
    92  	if deleted := idx.trie.Delete(patricia.Prefix(id)); !deleted {
    93  		return fmt.Errorf("no such id: '%s'", id)
    94  	}
    95  	return nil
    96  }
    97  
    98  // Get retrieves an ID from the TruncIndex. If there are multiple IDs
    99  // with the given prefix, an error is thrown.
   100  func (idx *TruncIndex) Get(s string) (string, error) {
   101  	if s == "" {
   102  		return "", ErrEmptyPrefix
   103  	}
   104  	var (
   105  		id string
   106  	)
   107  	subTreeVisitFunc := func(prefix patricia.Prefix, item patricia.Item) error {
   108  		if id != "" {
   109  			// we haven't found the ID if there are two or more IDs
   110  			id = ""
   111  			return ErrAmbiguousPrefix{prefix: s}
   112  		}
   113  		id = string(prefix)
   114  		return nil
   115  	}
   116  
   117  	idx.RLock()
   118  	defer idx.RUnlock()
   119  	if err := idx.trie.VisitSubtree(patricia.Prefix(s), subTreeVisitFunc); err != nil {
   120  		return "", err
   121  	}
   122  	if id != "" {
   123  		return id, nil
   124  	}
   125  	return "", ErrNotExist
   126  }
   127  
   128  // Iterate iterates over all stored IDs and passes each of them to the given
   129  // handler. Take care that the handler method does not call any public
   130  // method on truncindex as the internal locking is not reentrant/recursive
   131  // and will result in deadlock.
   132  func (idx *TruncIndex) Iterate(handler func(id string)) {
   133  	idx.Lock()
   134  	defer idx.Unlock()
   135  	idx.trie.Visit(func(prefix patricia.Prefix, item patricia.Item) error {
   136  		handler(string(prefix))
   137  		return nil
   138  	})
   139  }