Skip to content

Commit feae7de

Browse files
alexlovelltroyalexlovelltroy
authored
feat: enhance SMDClient with group membership caching and reverse indexing (#109)
* feat: enhance SMDClient with group membership caching and reverse indexing Signed-off-by: Alex Lovell-Troy <alovelltroy@lanl.gov> * refactor: replace if-else with switch for URL path handling in performance tests Signed-off-by: Alex Lovell-Troy <alovelltroy@lanl.gov> --------- Signed-off-by: Alex Lovell-Troy <alovelltroy@lanl.gov>
1 parent 573604c commit feae7de

3 files changed

Lines changed: 599 additions & 60 deletions

File tree

‎internal/smdclient/SMDclient.go‎

Lines changed: 80 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,14 @@ type SMDClient struct {
4646
tokenEndpoint string
4747
accessToken string
4848
nodes map[string]NodeMapping
49-
nodesMutex *sync.Mutex
49+
nodesMutex *sync.RWMutex
5050
nodes_last_update time.Time
5151
stopCacheRefresh chan struct{}
5252
stopOnce sync.Once
53+
// Reverse indexes for O(1) lookups
54+
ipToXname map[string]string
55+
macToXname map[string]string
56+
wgipToXname map[string]string
5357
}
5458

5559
type NodeInterface struct {
@@ -62,6 +66,7 @@ type NodeInterface struct {
6266
type NodeMapping struct {
6367
Xname string `json:"xname" yaml:"xname"`
6468
Interfaces []NodeInterface `json:"interfaces" yaml:"interfaces"`
69+
Groups []string `json:"groups" yaml:"groups"`
6570
}
6671

6772
// NewSMDClient creates a new SMDClient which connects to the SMD server at baseurl
@@ -105,10 +110,13 @@ func NewSMDClient(clusterName, baseurl, jwtURL, accessToken, certPath string, in
105110
smdBaseURL: baseurl,
106111
tokenEndpoint: jwtURL,
107112
accessToken: accessToken,
108-
nodesMutex: &sync.Mutex{},
113+
nodesMutex: &sync.RWMutex{},
109114
nodes_last_update: time.Now(),
110115
nodes: make(map[string]NodeMapping),
111116
stopCacheRefresh: make(chan struct{}),
117+
ipToXname: make(map[string]string),
118+
macToXname: make(map[string]string),
119+
wgipToXname: make(map[string]string),
112120
}
113121

114122
// Populate the cache initially
@@ -219,7 +227,7 @@ func (s *SMDClient) getSMD(ep string, smd interface{}) error {
219227
}
220228

221229
// PopulateNodes fetches the Ethernet interface data from the SMD server and populates the nodes map
222-
// with the corresponding node information, including MAC addresses, IP addresses, and descriptions.
230+
// with the corresponding node information, including MAC addresses, IP addresses, descriptions, and group membership.
223231
func (s *SMDClient) PopulateNodes() {
224232
s.nodesMutex.Lock()
225233
defer s.nodesMutex.Unlock()
@@ -273,44 +281,77 @@ func (s *SMDClient) PopulateNodes() {
273281
s.nodes[ep.CompID] = newNode
274282
}
275283
}
284+
285+
// Populate group membership for all nodes
286+
log.Debug().Msg("Fetching group membership for all nodes")
287+
for xname, node := range s.nodes {
288+
ml := new(sm.Membership)
289+
membershipEp := "/hsm/v2/memberships/" + xname
290+
if err := s.getSMD(membershipEp, ml); err != nil {
291+
log.Debug().Err(err).Msgf("Failed to get group membership for %s", xname)
292+
node.Groups = []string{} // Empty groups if fetch fails
293+
} else {
294+
node.Groups = ml.GroupLabels
295+
}
296+
s.nodes[xname] = node
297+
}
298+
299+
// Build reverse indexes for O(1) lookups
300+
log.Debug().Msg("Building reverse indexes")
301+
s.ipToXname = make(map[string]string)
302+
s.macToXname = make(map[string]string)
303+
s.wgipToXname = make(map[string]string)
304+
305+
for xname, node := range s.nodes {
306+
for _, iface := range node.Interfaces {
307+
if iface.IP != "" {
308+
s.ipToXname[strings.ToLower(iface.IP)] = xname
309+
}
310+
if iface.MAC != "" {
311+
s.macToXname[strings.ToLower(iface.MAC)] = xname
312+
}
313+
if iface.WGIP != "" {
314+
s.wgipToXname[strings.ToLower(iface.WGIP)] = xname
315+
}
316+
}
317+
}
318+
276319
s.nodes_last_update = time.Now()
277-
log.Debug().Msg("Nodes map populated")
320+
log.Debug().Msgf("Nodes map populated with %d nodes, %d IP mappings, %d MAC mappings",
321+
len(s.nodes), len(s.ipToXname), len(s.macToXname))
278322
}
279323

280324
// IDfromMAC returns the ID of the xname that has the MAC address
281325
func (s *SMDClient) IDfromMAC(mac string) (string, error) {
282-
s.nodesMutex.Lock()
283-
defer s.nodesMutex.Unlock()
326+
s.nodesMutex.RLock()
327+
defer s.nodesMutex.RUnlock()
284328

285-
for _, node := range s.nodes {
286-
for _, iface := range node.Interfaces {
287-
if strings.EqualFold(mac, iface.MAC) {
288-
return node.Xname, nil
289-
}
290-
}
329+
key := strings.ToLower(mac)
330+
if xname, found := s.macToXname[key]; found {
331+
return xname, nil
291332
}
292-
return "", errors.New("MAC " + mac + " not found for an xname in nodes")
333+
return "", fmt.Errorf("MAC %s not found for an xname in nodes", mac)
293334
}
294335

295336
// IDfromIP returns the ID of the xname that has the IP address
296337
func (s *SMDClient) IDfromIP(ipaddr string) (string, error) {
297-
s.nodesMutex.Lock()
298-
defer s.nodesMutex.Unlock()
338+
s.nodesMutex.RLock()
339+
defer s.nodesMutex.RUnlock()
299340

300-
for _, node := range s.nodes {
301-
for _, iface := range node.Interfaces {
302-
if strings.EqualFold(ipaddr, iface.IP) || strings.EqualFold(ipaddr, iface.WGIP) {
303-
return node.Xname, nil
304-
}
305-
}
341+
key := strings.ToLower(ipaddr)
342+
if xname, found := s.ipToXname[key]; found {
343+
return xname, nil
306344
}
307-
return "", errors.New("IP address " + ipaddr + " not found for an xname in nodes")
345+
if xname, found := s.wgipToXname[key]; found {
346+
return xname, nil
347+
}
348+
return "", fmt.Errorf("IP address %s not found for an xname in nodes", ipaddr)
308349
}
309350

310351
// IPfromID returns the IP address of the xname with the given ID
311352
func (s *SMDClient) IPfromID(id string) (string, error) {
312-
s.nodesMutex.Lock()
313-
defer s.nodesMutex.Unlock()
353+
s.nodesMutex.RLock()
354+
defer s.nodesMutex.RUnlock()
314355
if node, found := s.nodes[id]; found {
315356
if node.Interfaces != nil {
316357
if len(node.Interfaces) > 0 {
@@ -323,8 +364,8 @@ func (s *SMDClient) IPfromID(id string) (string, error) {
323364
}
324365

325366
func (s *SMDClient) MACfromID(id string) (string, error) {
326-
s.nodesMutex.Lock()
327-
defer s.nodesMutex.Unlock()
367+
s.nodesMutex.RLock()
368+
defer s.nodesMutex.RUnlock()
328369
if node, found := s.nodes[id]; found {
329370
if node.Interfaces != nil {
330371
if len(node.Interfaces) > 0 {
@@ -343,13 +384,15 @@ func (s *SMDClient) GroupMembership(id string) ([]string, error) {
343384
log.Err(err).Msg("failed to get group membership")
344385
return []string{}, err
345386
}
346-
ml := new(sm.Membership)
347-
ep := "/hsm/v2/memberships/" + id
348-
err := s.getSMD(ep, ml)
349-
if err != nil {
350-
return nil, err
387+
388+
s.nodesMutex.RLock()
389+
defer s.nodesMutex.RUnlock()
390+
391+
if node, found := s.nodes[id]; found {
392+
return node.Groups, nil
351393
}
352-
return ml.GroupLabels, nil
394+
395+
return []string{}, fmt.Errorf("node %s not found in cache", id)
353396
}
354397

355398
func (s *SMDClient) ComponentInformation(id string) (base.Component, error) {
@@ -372,6 +415,9 @@ func (s *SMDClient) AddWGIP(id string, wgip string) error {
372415
if node.Interfaces != nil {
373416
if len(node.Interfaces) > 0 {
374417
node.Interfaces[0].WGIP = wgip
418+
s.nodes[id] = node
419+
// Update reverse index
420+
s.wgipToXname[strings.ToLower(wgip)] = id
375421
return nil
376422
}
377423
return errors.New("no interfaces found for ID " + id)
@@ -381,8 +427,8 @@ func (s *SMDClient) AddWGIP(id string, wgip string) error {
381427
}
382428

383429
func (s *SMDClient) WGIPfromID(id string) (string, error) {
384-
s.nodesMutex.Lock()
385-
defer s.nodesMutex.Unlock()
430+
s.nodesMutex.RLock()
431+
defer s.nodesMutex.RUnlock()
386432
if node, found := s.nodes[id]; found {
387433
if node.Interfaces != nil {
388434
if len(node.Interfaces) > 0 {

0 commit comments

Comments
 (0)