all: imp code

This commit is contained in:
Stanislav Chzhen 2024-08-20 16:45:43 +03:00
parent 243123a404
commit 6567cd330c
4 changed files with 31 additions and 23 deletions

View File

@ -121,6 +121,7 @@ func (r *Runtime) Info() (cs Source, host string) {
// SetInfo sets a host as a client information from the cs. // SetInfo sets a host as a client information from the cs.
func (r *Runtime) SetInfo(cs Source, hosts []string) { 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] == "" { if len(hosts) == 1 && hosts[0] == "" {
hosts = []string{} hosts = []string{}
} }

View File

@ -237,21 +237,21 @@ func (s *Storage) ClientRuntime(ip netip.Addr) (rc *Runtime) {
// UpdateRuntime updates the stored runtime client with information from rc. If // 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. // 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() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
s.updateRuntimeLocked(rc) return s.updateRuntimeLocked(rc)
} }
// updateRuntimeLocked updates the stored runtime client with information from // updateRuntimeLocked updates the stored runtime client with information from
// rc. rc must not be nil. Storage.mu is expected to be locked. // 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) stored := s.runtimeIndex.Client(rc.ip)
if stored == nil { if stored == nil {
s.runtimeIndex.Add(rc.Clone()) s.runtimeIndex.Add(rc.Clone())
return return true
} }
if rc.whois != nil { if rc.whois != nil {
@ -273,11 +273,13 @@ func (s *Storage) updateRuntimeLocked(rc *Runtime) {
if rc.hostsFile != nil { if rc.hostsFile != nil {
stored.hostsFile = slices.Clone(rc.hostsFile) stored.hostsFile = slices.Clone(rc.hostsFile)
} }
return false
} }
// BatchUpdateBySource updates the stored runtime clients information from the // BatchUpdateBySource updates the stored runtime clients information from the
// specified source. // specified source and returns the number of added and removed clients.
func (s *Storage) BatchUpdateBySource(src Source, rcs []*Runtime) { func (s *Storage) BatchUpdateBySource(src Source, rcs []*Runtime) (added, removed int) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
@ -286,14 +288,19 @@ func (s *Storage) BatchUpdateBySource(src Source, rcs []*Runtime) {
} }
for _, rc := range rcs { for _, rc := range rcs {
s.updateRuntimeLocked(rc) if s.updateRuntimeLocked(rc) {
added++
}
} }
for ip, rc := range s.runtimeIndex.index { for ip, rc := range s.runtimeIndex.index {
if rc.isEmpty() { if rc.isEmpty() {
delete(s.runtimeIndex.index, ip) delete(s.runtimeIndex.index, ip)
removed++
} }
} }
return added, removed
} }
// SizeRuntime returns the number of the runtime clients. // SizeRuntime returns the number of the runtime clients.

View File

@ -659,7 +659,7 @@ func TestStorage_UpdateRuntime(t *testing.T) {
added := client.NewRuntime(ip) added := client.NewRuntime(ip)
added.SetInfo(client.SourceARP, []string{addedARP}) added.SetInfo(client.SourceARP, []string{addedARP})
s.UpdateRuntime(added) require.True(t, s.UpdateRuntime(added))
require.Equal(t, 1, s.SizeRuntime()) require.Equal(t, 1, s.SizeRuntime())
got := s.ClientRuntime(ip) got := s.ClientRuntime(ip)
@ -672,7 +672,7 @@ func TestStorage_UpdateRuntime(t *testing.T) {
added := client.NewRuntime(ip2) added := client.NewRuntime(ip2)
added.SetInfo(client.SourceARP, []string{addedSecondARP}) added.SetInfo(client.SourceARP, []string{addedSecondARP})
s.UpdateRuntime(added) require.True(t, s.UpdateRuntime(added))
require.Equal(t, 2, s.SizeRuntime()) require.Equal(t, 2, s.SizeRuntime())
got := s.ClientRuntime(ip2) got := s.ClientRuntime(ip2)
@ -682,7 +682,7 @@ func TestStorage_UpdateRuntime(t *testing.T) {
}) })
t.Run("update_first_client", func(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) got := s.ClientRuntime(ip)
require.Equal(t, 2, s.SizeRuntime()) require.Equal(t, 2, s.SizeRuntime())
@ -744,7 +744,9 @@ func TestStorage_BatchUpdateBySource(t *testing.T) {
}) })
t.Run("populate_storage_with_first_clients", func(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()) require.Equal(t, len(firstClients), s.SizeRuntime())
rc := s.ClientRuntime(cliFirstIP1) rc := s.ClientRuntime(cliFirstIP1)
@ -754,7 +756,9 @@ func TestStorage_BatchUpdateBySource(t *testing.T) {
}) })
t.Run("update_storage", func(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()) require.Equal(t, len(updatedClients), s.SizeRuntime())
rc := s.ClientRuntime(cliUpdatedIP3) rc := s.ClientRuntime(cliUpdatedIP3)
@ -767,7 +771,9 @@ func TestStorage_BatchUpdateBySource(t *testing.T) {
}) })
t.Run("remove_all", func(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()) require.Equal(t, 0, s.SizeRuntime())
}) })
} }

View File

@ -642,7 +642,6 @@ func (clients *clientsContainer) addFromHostsFile(hosts *hostsfile.DefaultStorag
clients.lock.Lock() clients.lock.Lock()
defer clients.lock.Unlock() defer clients.lock.Unlock()
added := 0
var rcs []*client.Runtime var rcs []*client.Runtime
hosts.RangeNames(func(addr netip.Addr, names []string) (cont bool) { hosts.RangeNames(func(addr netip.Addr, names []string) (cont bool) {
// Only the first name of the first record is considered a canonical // Only the first name of the first record is considered a canonical
@ -652,15 +651,13 @@ func (clients *clientsContainer) addFromHostsFile(hosts *hostsfile.DefaultStorag
rc := client.NewRuntime(addr) rc := client.NewRuntime(addr)
rc.SetInfo(client.SourceHostsFile, []string{names[0]}) rc.SetInfo(client.SourceHostsFile, []string{names[0]})
added++
rcs = append(rcs, rc) rcs = append(rcs, rc)
return true return true
}) })
clients.storage.BatchUpdateBySource(client.SourceHostsFile, rcs) added, removed := clients.storage.BatchUpdateBySource(client.SourceHostsFile, rcs)
log.Debug("clients: added %d, removed %d client aliases from system hosts file", added, removed)
log.Debug("clients: added %d client aliases from system hosts file", added)
} }
// addFromSystemARP adds the IP-hostname pairings from the output of the arp -a // addFromSystemARP adds the IP-hostname pairings from the output of the arp -a
@ -684,19 +681,16 @@ func (clients *clientsContainer) addFromSystemARP() {
clients.lock.Lock() clients.lock.Lock()
defer clients.lock.Unlock() defer clients.lock.Unlock()
added := 0
var rcs []*client.Runtime var rcs []*client.Runtime
for _, n := range ns { for _, n := range ns {
rc := client.NewRuntime(n.IP) rc := client.NewRuntime(n.IP)
rc.SetInfo(client.SourceARP, []string{n.Name}) rc.SetInfo(client.SourceARP, []string{n.Name})
added++
rcs = append(rcs, rc) rcs = append(rcs, rc)
} }
clients.storage.BatchUpdateBySource(client.SourceARP, rcs) added, removed := clients.storage.BatchUpdateBySource(client.SourceARP, rcs)
log.Debug("clients: added %d, removed %d client aliases from arp neighborhood", added, removed)
log.Debug("clients: added %d client aliases from arp neighborhood", added)
} }
// close gracefully closes all the client-specific upstream configurations of // close gracefully closes all the client-specific upstream configurations of