truncindex.go 3.6 KB

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