Skip to content

Commit

Permalink
all: imp code
Browse files Browse the repository at this point in the history
  • Loading branch information
schzhn committed Aug 20, 2024
1 parent 243123a commit 6567cd3
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 23 deletions.
1 change: 1 addition & 0 deletions internal/client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,7 @@ func (r *Runtime) Info() (cs Source, host string) {

// SetInfo sets a host as a client information from the cs.
func (r *Runtime) SetInfo(cs Source, hosts []string) {
// TODO(s.chzhen): Use contract where hosts must contain non-empty host.
if len(hosts) == 1 && hosts[0] == "" {
hosts = []string{}
}
Expand Down
21 changes: 14 additions & 7 deletions internal/client/storage.go
Original file line number Diff line number Diff line change
Expand Up @@ -237,21 +237,21 @@ func (s *Storage) ClientRuntime(ip netip.Addr) (rc *Runtime) {

// UpdateRuntime updates the stored runtime client with information from rc. If
// no such client exists, saves the copy of rc in storage. rc must not be nil.
func (s *Storage) UpdateRuntime(rc *Runtime) {
func (s *Storage) UpdateRuntime(rc *Runtime) (added bool) {
s.mu.Lock()
defer s.mu.Unlock()

s.updateRuntimeLocked(rc)
return s.updateRuntimeLocked(rc)
}

// updateRuntimeLocked updates the stored runtime client with information from
// rc. rc must not be nil. Storage.mu is expected to be locked.
func (s *Storage) updateRuntimeLocked(rc *Runtime) {
func (s *Storage) updateRuntimeLocked(rc *Runtime) (added bool) {
stored := s.runtimeIndex.Client(rc.ip)
if stored == nil {
s.runtimeIndex.Add(rc.Clone())

return
return true
}

if rc.whois != nil {
Expand All @@ -273,11 +273,13 @@ func (s *Storage) updateRuntimeLocked(rc *Runtime) {
if rc.hostsFile != nil {
stored.hostsFile = slices.Clone(rc.hostsFile)
}

return false
}

// BatchUpdateBySource updates the stored runtime clients information from the
// specified source.
func (s *Storage) BatchUpdateBySource(src Source, rcs []*Runtime) {
// specified source and returns the number of added and removed clients.
func (s *Storage) BatchUpdateBySource(src Source, rcs []*Runtime) (added, removed int) {
s.mu.Lock()
defer s.mu.Unlock()

Expand All @@ -286,14 +288,19 @@ func (s *Storage) BatchUpdateBySource(src Source, rcs []*Runtime) {
}

for _, rc := range rcs {
s.updateRuntimeLocked(rc)
if s.updateRuntimeLocked(rc) {
added++
}
}

for ip, rc := range s.runtimeIndex.index {
if rc.isEmpty() {
delete(s.runtimeIndex.index, ip)
removed++
}
}

return added, removed
}

// SizeRuntime returns the number of the runtime clients.
Expand Down
18 changes: 12 additions & 6 deletions internal/client/storage_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -659,7 +659,7 @@ func TestStorage_UpdateRuntime(t *testing.T) {
added := client.NewRuntime(ip)
added.SetInfo(client.SourceARP, []string{addedARP})

s.UpdateRuntime(added)
require.True(t, s.UpdateRuntime(added))
require.Equal(t, 1, s.SizeRuntime())

got := s.ClientRuntime(ip)
Expand All @@ -672,7 +672,7 @@ func TestStorage_UpdateRuntime(t *testing.T) {
added := client.NewRuntime(ip2)
added.SetInfo(client.SourceARP, []string{addedSecondARP})

s.UpdateRuntime(added)
require.True(t, s.UpdateRuntime(added))
require.Equal(t, 2, s.SizeRuntime())

got := s.ClientRuntime(ip2)
Expand All @@ -682,7 +682,7 @@ func TestStorage_UpdateRuntime(t *testing.T) {
})

t.Run("update_first_client", func(t *testing.T) {
s.UpdateRuntime(updated)
require.False(t, s.UpdateRuntime(updated))
got := s.ClientRuntime(ip)
require.Equal(t, 2, s.SizeRuntime())

Expand Down Expand Up @@ -744,7 +744,9 @@ func TestStorage_BatchUpdateBySource(t *testing.T) {
})

t.Run("populate_storage_with_first_clients", func(t *testing.T) {
s.BatchUpdateBySource(defSrc, firstClients)
added, removed := s.BatchUpdateBySource(defSrc, firstClients)
require.Equal(t, len(firstClients), added)
require.Equal(t, 0, removed)
require.Equal(t, len(firstClients), s.SizeRuntime())

rc := s.ClientRuntime(cliFirstIP1)
Expand All @@ -754,7 +756,9 @@ func TestStorage_BatchUpdateBySource(t *testing.T) {
})

t.Run("update_storage", func(t *testing.T) {
s.BatchUpdateBySource(defSrc, updatedClients)
added, removed := s.BatchUpdateBySource(defSrc, updatedClients)
require.Equal(t, len(updatedClients), added)
require.Equal(t, len(firstClients), removed)
require.Equal(t, len(updatedClients), s.SizeRuntime())

rc := s.ClientRuntime(cliUpdatedIP3)
Expand All @@ -767,7 +771,9 @@ func TestStorage_BatchUpdateBySource(t *testing.T) {
})

t.Run("remove_all", func(t *testing.T) {
s.BatchUpdateBySource(defSrc, []*client.Runtime{})
added, removed := s.BatchUpdateBySource(defSrc, []*client.Runtime{})
require.Equal(t, 0, added)
require.Equal(t, len(updatedClients), removed)
require.Equal(t, 0, s.SizeRuntime())
})
}
14 changes: 4 additions & 10 deletions internal/home/clients.go
Original file line number Diff line number Diff line change
Expand Up @@ -642,7 +642,6 @@ func (clients *clientsContainer) addFromHostsFile(hosts *hostsfile.DefaultStorag
clients.lock.Lock()
defer clients.lock.Unlock()

added := 0
var rcs []*client.Runtime
hosts.RangeNames(func(addr netip.Addr, names []string) (cont bool) {
// Only the first name of the first record is considered a canonical
Expand All @@ -652,15 +651,13 @@ func (clients *clientsContainer) addFromHostsFile(hosts *hostsfile.DefaultStorag
rc := client.NewRuntime(addr)
rc.SetInfo(client.SourceHostsFile, []string{names[0]})

added++
rcs = append(rcs, rc)

return true
})

clients.storage.BatchUpdateBySource(client.SourceHostsFile, rcs)

log.Debug("clients: added %d client aliases from system hosts file", added)
added, removed := clients.storage.BatchUpdateBySource(client.SourceHostsFile, rcs)
log.Debug("clients: added %d, removed %d client aliases from system hosts file", added, removed)
}

// addFromSystemARP adds the IP-hostname pairings from the output of the arp -a
Expand All @@ -684,19 +681,16 @@ func (clients *clientsContainer) addFromSystemARP() {
clients.lock.Lock()
defer clients.lock.Unlock()

added := 0
var rcs []*client.Runtime
for _, n := range ns {
rc := client.NewRuntime(n.IP)
rc.SetInfo(client.SourceARP, []string{n.Name})

added++
rcs = append(rcs, rc)
}

clients.storage.BatchUpdateBySource(client.SourceARP, rcs)

log.Debug("clients: added %d client aliases from arp neighborhood", added)
added, removed := clients.storage.BatchUpdateBySource(client.SourceARP, rcs)
log.Debug("clients: added %d, removed %d client aliases from arp neighborhood", added, removed)
}

// close gracefully closes all the client-specific upstream configurations of
Expand Down

0 comments on commit 6567cd3

Please sign in to comment.