mirror of
https://github.com/therootcompany/golib.git
synced 2026-11-04 06:55:00 +00:00
feat(iplist): add refreshable Source
This commit is contained in:
parent
1cb0235d9d
commit
869086ce0c
143
net/iplist/source.go
Normal file
143
net/iplist/source.go
Normal file
@ -0,0 +1,143 @@
|
||||
package iplist
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
const DefaultRefreshInterval = time.Hour + 57*time.Minute + 13*time.Second
|
||||
|
||||
type SourceConfig struct {
|
||||
Source string
|
||||
CacheDir string
|
||||
HTTPClient *http.Client
|
||||
RefreshInterval time.Duration
|
||||
Optional bool
|
||||
OnRefresh func(SourceEvent)
|
||||
}
|
||||
|
||||
type SourceEvent struct {
|
||||
Entries uint64
|
||||
Err error
|
||||
}
|
||||
|
||||
type sourceSnapshot struct {
|
||||
entries []string
|
||||
}
|
||||
|
||||
type Source struct {
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
config SourceConfig
|
||||
current atomic.Pointer[sourceSnapshot]
|
||||
observers []func(SourceEvent)
|
||||
observerMu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewSource(ctx context.Context, config SourceConfig) (*Source, error) {
|
||||
if config.RefreshInterval <= 0 {
|
||||
config.RefreshInterval = DefaultRefreshInterval
|
||||
}
|
||||
sourceCtx, cancel := context.WithCancel(ctx)
|
||||
s := &Source{ctx: sourceCtx, cancel: cancel, config: config}
|
||||
entries, err := s.load()
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, err
|
||||
}
|
||||
s.store(entries)
|
||||
go s.refreshLoop()
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *Source) Entries() []string {
|
||||
current := s.current.Load()
|
||||
if current == nil {
|
||||
return nil
|
||||
}
|
||||
return append([]string(nil), current.entries...)
|
||||
}
|
||||
|
||||
func (s *Source) Refresh() error {
|
||||
entries, err := s.load()
|
||||
if err != nil {
|
||||
s.emit(SourceEvent{Entries: uint64(len(s.Entries())), Err: err})
|
||||
return err
|
||||
}
|
||||
s.store(entries)
|
||||
s.emit(SourceEvent{Entries: uint64(len(entries))})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Source) Subscribe(observer func(SourceEvent)) {
|
||||
if observer == nil {
|
||||
return
|
||||
}
|
||||
s.observerMu.Lock()
|
||||
s.observers = append(s.observers, observer)
|
||||
s.observerMu.Unlock()
|
||||
}
|
||||
|
||||
func (s *Source) emit(event SourceEvent) {
|
||||
if s.config.OnRefresh != nil {
|
||||
s.config.OnRefresh(event)
|
||||
}
|
||||
s.observerMu.RLock()
|
||||
observers := append([]func(SourceEvent){}, s.observers...)
|
||||
s.observerMu.RUnlock()
|
||||
for _, observer := range observers {
|
||||
observer(event)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Source) Close() error {
|
||||
if s != nil && s.cancel != nil {
|
||||
s.cancel()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Source) load() ([]string, error) {
|
||||
if s.config.Source == "" {
|
||||
if s.config.Optional {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, errors.New("source is required")
|
||||
}
|
||||
if s.config.Optional {
|
||||
parsed, parseErr := url.Parse(s.config.Source)
|
||||
isURL := parseErr == nil && (parsed.Scheme == "http" || parsed.Scheme == "https")
|
||||
if !isURL {
|
||||
if _, err := os.Stat(s.config.Source); os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
} else if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
return Load(s.ctx, s.config.Source, s.config.CacheDir, s.config.HTTPClient)
|
||||
}
|
||||
|
||||
func (s *Source) store(entries []string) {
|
||||
copyEntries := append([]string(nil), entries...)
|
||||
s.current.Store(&sourceSnapshot{entries: copyEntries})
|
||||
}
|
||||
|
||||
func (s *Source) refreshLoop() {
|
||||
ticker := time.NewTicker(s.config.RefreshInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
_ = s.Refresh()
|
||||
}
|
||||
}
|
||||
}
|
||||
46
net/iplist/source_test.go
Normal file
46
net/iplist/source_test.go
Normal file
@ -0,0 +1,46 @@
|
||||
package iplist
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSourceKeepsLastGoodEntriesAfterFailedRefresh(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "list.tsv")
|
||||
if err := os.WriteFile(path, []byte("192.0.2.1\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
source, err := NewSource(t.Context(), SourceConfig{Source: path})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = source.Close() }()
|
||||
|
||||
if got := source.Entries(); len(got) != 1 || got[0] != "192.0.2.1" {
|
||||
t.Fatalf("initial entries = %#v", got)
|
||||
}
|
||||
if err := os.Remove(path); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := source.Refresh(); err == nil {
|
||||
t.Fatal("invalid refresh unexpectedly succeeded")
|
||||
}
|
||||
if got := source.Entries(); len(got) != 1 || got[0] != "192.0.2.1" {
|
||||
t.Fatalf("entries after failed refresh = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSourceOptionalMissingFileIsEmpty(t *testing.T) {
|
||||
source, err := NewSource(t.Context(), SourceConfig{
|
||||
Source: filepath.Join(t.TempDir(), "missing.tsv"),
|
||||
Optional: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = source.Close() }()
|
||||
if got := source.Entries(); len(got) != 0 {
|
||||
t.Fatalf("entries = %#v", got)
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user