github.com/hernad/nomad@v1.6.112/nomad/state/state_store_node_pools.go (about) 1 // Copyright (c) HashiCorp, Inc. 2 // SPDX-License-Identifier: MPL-2.0 3 4 package state 5 6 import ( 7 "fmt" 8 9 memdb "github.com/hashicorp/go-memdb" 10 "github.com/hernad/nomad/nomad/structs" 11 ) 12 13 // nodePoolInit creates the built-in node pools that should always be present 14 // in the cluster. 15 func (s *StateStore) nodePoolInit() error { 16 allNodePool := &structs.NodePool{ 17 Name: structs.NodePoolAll, 18 Description: structs.NodePoolAllDescription, 19 } 20 21 defaultNodePool := &structs.NodePool{ 22 Name: structs.NodePoolDefault, 23 Description: structs.NodePoolDefaultDescription, 24 } 25 26 return s.UpsertNodePools( 27 structs.SystemInitializationType, 28 1, 29 []*structs.NodePool{allNodePool, defaultNodePool}, 30 ) 31 } 32 33 // NodePools returns an iterator over all node pools. 34 func (s *StateStore) NodePools(ws memdb.WatchSet, sort SortOption) (memdb.ResultIterator, error) { 35 txn := s.db.ReadTxn() 36 37 var iter memdb.ResultIterator 38 var err error 39 40 switch sort { 41 case SortReverse: 42 iter, err = txn.GetReverse(TableNodePools, "id") 43 default: 44 iter, err = txn.Get(TableNodePools, "id") 45 } 46 if err != nil { 47 return nil, fmt.Errorf("node pools lookup failed: %w", err) 48 } 49 50 ws.Add(iter.WatchCh()) 51 return iter, nil 52 } 53 54 // NodePoolByName returns the node pool that matches the given name or nil if 55 // there is no match. 56 func (s *StateStore) NodePoolByName(ws memdb.WatchSet, name string) (*structs.NodePool, error) { 57 txn := s.db.ReadTxn() 58 return s.nodePoolByNameTxn(txn, ws, name) 59 } 60 61 func (s *StateStore) nodePoolByNameTxn(txn *txn, ws memdb.WatchSet, name string) (*structs.NodePool, error) { 62 watchCh, existing, err := txn.FirstWatch(TableNodePools, "id", name) 63 if err != nil { 64 return nil, fmt.Errorf("node pool lookup failed: %w", err) 65 } 66 ws.Add(watchCh) 67 68 if existing == nil { 69 return nil, nil 70 } 71 72 return existing.(*structs.NodePool), nil 73 } 74 75 // NodePoolsByNamePrefix returns an interator over all node pools that match 76 // the given name prefix. 77 func (s *StateStore) NodePoolsByNamePrefix(ws memdb.WatchSet, namePrefix string, sort SortOption) (memdb.ResultIterator, error) { 78 txn := s.db.ReadTxn() 79 80 var iter memdb.ResultIterator 81 var err error 82 83 switch sort { 84 case SortReverse: 85 iter, err = txn.GetReverse(TableNodePools, "id_prefix", namePrefix) 86 default: 87 iter, err = txn.Get(TableNodePools, "id_prefix", namePrefix) 88 } 89 if err != nil { 90 return nil, fmt.Errorf("node pools prefix lookup failed: %w", err) 91 } 92 93 ws.Add(iter.WatchCh()) 94 return iter, nil 95 } 96 97 // nodePoolExists returs true if a node pool with the give name exists. 98 func (s *StateStore) nodePoolExists(txn *txn, pool string) (bool, error) { 99 existing, err := txn.First(TableNodePools, "id", pool) 100 return existing != nil, err 101 } 102 103 // UpsertNodePools inserts or updates the given set of node pools. 104 func (s *StateStore) UpsertNodePools(msgType structs.MessageType, index uint64, pools []*structs.NodePool) error { 105 txn := s.db.WriteTxnMsgT(msgType, index) 106 defer txn.Abort() 107 108 for _, pool := range pools { 109 if err := s.upsertNodePoolTxn(txn, index, pool); err != nil { 110 return err 111 } 112 } 113 114 if err := txn.Insert("index", &IndexEntry{TableNodePools, index}); err != nil { 115 return fmt.Errorf("index update failed: %w", err) 116 } 117 118 return txn.Commit() 119 } 120 121 func (s *StateStore) upsertNodePoolTxn(txn *txn, index uint64, pool *structs.NodePool) error { 122 if pool == nil { 123 return nil 124 } 125 126 existing, err := txn.First(TableNodePools, "id", pool.Name) 127 if err != nil { 128 return fmt.Errorf("node pool lookup failed: %w", err) 129 } 130 131 if existing != nil { 132 // Prevent changes to built-in node pools. 133 if pool.IsBuiltIn() { 134 return fmt.Errorf("modifying node pool %q is not allowed", pool.Name) 135 } 136 137 exist := existing.(*structs.NodePool) 138 pool.CreateIndex = exist.CreateIndex 139 pool.ModifyIndex = index 140 } else { 141 pool.CreateIndex = index 142 pool.ModifyIndex = index 143 } 144 145 if err := txn.Insert(TableNodePools, pool); err != nil { 146 return fmt.Errorf("node pool insert failed: %w", err) 147 } 148 149 return nil 150 } 151 152 // fetchOrCreateNodePoolTxn returns an existing node pool with the given name 153 // or creates a new one if it doesn't exist. 154 func (s *StateStore) fetchOrCreateNodePoolTxn(txn *txn, index uint64, name string) (*structs.NodePool, error) { 155 pool, err := s.nodePoolByNameTxn(txn, nil, name) 156 if err != nil { 157 return nil, err 158 } 159 160 if pool == nil { 161 pool = &structs.NodePool{Name: name} 162 err = s.upsertNodePoolTxn(txn, index, pool) 163 if err != nil { 164 return nil, err 165 } 166 } 167 168 return pool, nil 169 } 170 171 // DeleteNodePools removes the given set of node pools. 172 func (s *StateStore) DeleteNodePools(msgType structs.MessageType, index uint64, names []string) error { 173 txn := s.db.WriteTxnMsgT(msgType, index) 174 defer txn.Abort() 175 176 for _, n := range names { 177 if err := s.deleteNodePoolTxn(txn, index, n); err != nil { 178 return err 179 } 180 } 181 182 // Update index table. 183 if err := txn.Insert("index", &IndexEntry{TableNodePools, index}); err != nil { 184 return fmt.Errorf("index update failed: %w", err) 185 } 186 187 return txn.Commit() 188 } 189 190 func (s *StateStore) deleteNodePoolTxn(txn *txn, index uint64, name string) error { 191 // Check if node pool exists. 192 existing, err := txn.First(TableNodePools, "id", name) 193 if err != nil { 194 return fmt.Errorf("node pool lookup failed: %w", err) 195 } 196 if existing == nil { 197 return fmt.Errorf("node pool %s not found", name) 198 } 199 200 pool := existing.(*structs.NodePool) 201 202 // Prevent deletion of built-in node pools. 203 if pool.IsBuiltIn() { 204 return fmt.Errorf("deleting node pool %q is not allowed", pool.Name) 205 } 206 207 // Delete node pool. 208 if err := txn.Delete(TableNodePools, pool); err != nil { 209 return fmt.Errorf("node pool deletion failed: %w", err) 210 } 211 212 return nil 213 }