Better concurrency handling

Author: myleshorton

cache.go

@@ -42,7 +42,7 @@ func (f *fronted) prepopulateFronts(cacheFile string) {
 	now := time.Now()
 
 	// update last succeeded status of masquerades based on cached values
-	for _, fr := range f.fronts {
+	for _, fr := range f.fronts.fronts {
 		for _, cf := range cachedFronts {
 			sameFront := cf.ProviderID == fr.getProviderID() && cf.Domain == fr.getDomain() && cf.IpAddress == fr.getIpAddress()
 			cachedValueFresh := now.Sub(fr.lastSucceeded()) < f.maxAllowedCachedAge

cache_test.go

@@ -29,7 +29,7 @@ func TestCaching(t *testing.T) {
 	log.Debug("Creating fronted")
 	makeFronted := func() *fronted {
 		f := &fronted{
-			fronts:              make(sortedFronts, 0, 1000),
+			fronts:              newThreadSafeFronts(1000),
 			maxAllowedCachedAge: 250 * time.Millisecond,
 			maxCacheSize:        4,
 			cacheSaveInterval:   50 * time.Millisecond,
@@ -51,7 +51,7 @@ func TestCaching(t *testing.T) {
 	f := makeFronted()
 
 	log.Debug("Adding fronts")
-	f.fronts = append(f.fronts, mb, mc, md)
+	f.fronts.fronts = append(f.fronts.fronts, mb, mc, md)
 
 	readCached := func() []*front {
 		log.Debug("Reading cached fronts")

front.go

@@ -360,6 +360,45 @@ func NewStatusCodeValidator(reject []int) ResponseValidator {
 	}
 }
 
+type threadSafeFronts struct {
+	fronts sortedFronts
+	mx     sync.RWMutex
+}
+
+func newThreadSafeFronts(size int) *threadSafeFronts {
+	return &threadSafeFronts{
+		fronts: make(sortedFronts, 0, size),
+		mx:     sync.RWMutex{},
+	}
+}
+
+func (tsf *threadSafeFronts) sortedCopy() sortedFronts {
+	tsf.mx.RLock()
+	defer tsf.mx.RUnlock()
+	c := make(sortedFronts, len(tsf.fronts))
+	copy(c, tsf.fronts)
+	sort.Sort(c)
+	return c
+}
+
+func (tsf *threadSafeFronts) addFronts(newFronts ...Front) {
+	tsf.mx.Lock()
+	defer tsf.mx.Unlock()
+	tsf.fronts = append(tsf.fronts, newFronts...)
+}
+
+func (tsf *threadSafeFronts) frontSize() int {
+	tsf.mx.RLock()
+	defer tsf.mx.RUnlock()
+	return len(tsf.fronts)
+}
+
+func (tsf *threadSafeFronts) frontAt(i int) Front {
+	tsf.mx.RLock()
+	defer tsf.mx.RUnlock()
+	return tsf.fronts[i]
+}
+
 // slice of masquerade sorted by last vetted time
 type sortedFronts []Front
 
@@ -375,13 +414,6 @@ func (m sortedFronts) Less(i, j int) bool {
 	}
 }
 
-func (m sortedFronts) sortedCopy() sortedFronts {
-	c := make(sortedFronts, len(m))
-	copy(c, m)
-	sort.Sort(c)
-	return c
-}
-
 func (fr *front) markCacheDirty() {
 	select {
 	case fr.cacheDirty <- nil:

fronted.go

@@ -49,7 +49,7 @@ var (
 // an implementation of http.RoundTripper for the convenience of callers.
 type fronted struct {
 	certPool            atomic.Value
-	fronts              sortedFronts
+	fronts              *threadSafeFronts
 	maxAllowedCachedAge time.Duration
 	maxCacheSize        int
 	cacheFile           string
@@ -102,7 +102,7 @@ func NewFronted(options ...Option) Fronted {
 
 	f := &fronted{
 		certPool:            atomic.Value{},
-		fronts:              make(sortedFronts, 0),
+		fronts:              newThreadSafeFronts(0),
 		maxAllowedCachedAge: defaultMaxAllowedCachedAge,
 		maxCacheSize:        defaultMaxCacheSize,
 		cacheSaveInterval:   defaultCacheSaveInterval,
@@ -260,7 +260,7 @@ func (f *fronted) onNewFronts(pool *x509.CertPool, providers map[string]*Provide
 	}
 	providersCopy := copyProviders(providers, f.countryCode)
 	f.addProviders(providersCopy)
-	f.addFronts(loadFronts(providersCopy, f.cacheDirty))
+	f.fronts.addFronts(loadFronts(providersCopy, f.cacheDirty)...)
 	f.certPool.Store(pool)
 
 	// The goroutine for finding working fronts runs forever, so only start it once.
@@ -320,8 +320,8 @@ func (f *fronted) tryAllFronts() {
 	pool := pond.NewPool(40)
 
 	// Submit all fronts to the worker pool.
-	for i := range f.frontSize() {
-		m := f.frontAt(i)
+	for i := range f.fronts.frontSize() {
+		m := f.fronts.frontAt(i)
 		pool.Submit(func() {
 			if f.isStopped() {
 				return
@@ -348,18 +348,6 @@ func (f *fronted) hasEnoughWorkingFronts() bool {
 	return len(f.frontsCh) >= 4
 }
 
-func (f *fronted) frontSize() int {
-	f.frontsMu.RLock()
-	defer f.frontsMu.RUnlock()
-	return len(f.fronts)
-}
-
-func (f *fronted) frontAt(i int) Front {
-	f.frontsMu.RLock()
-	defer f.frontsMu.RUnlock()
-	return f.fronts[i]
-}
-
 func (f *fronted) vetFront(fr Front) bool {
 	conn, err := f.dialFront(fr)
 	if err != nil {
@@ -571,7 +559,7 @@ func copyProviders(providers map[string]*Provider, countryCode string) map[strin
 	return providersCopy
 }
 
-func loadFronts(providers map[string]*Provider, cacheDirty chan interface{}) sortedFronts {
+func loadFronts(providers map[string]*Provider, cacheDirty chan interface{}) []Front {
 	log.Debugf("Loading candidates for %d providers", len(providers))
 	defer log.Debug("Finished loading candidates")
 
@@ -581,7 +569,7 @@ func loadFronts(providers map[string]*Provider, cacheDirty chan interface{}) sor
 		size += len(p.Masquerades)
 	}
 
-	fronts := make(sortedFronts, size)
+	fronts := make([]Front, size)
 
 	// Note that map iteration order is random, so the order of the providers is automatically randomized.
 	index := 0
@@ -616,13 +604,6 @@ func (f *fronted) addProviders(providers map[string]*Provider) {
 	}
 }
 
-func (f *fronted) addFronts(fronts sortedFronts) {
-	// Add new masquerades to the existing masquerades slice, but add them at the beginning.
-	f.frontsMu.Lock()
-	defer f.frontsMu.Unlock()
-	f.fronts = append(fronts, f.fronts.sortedCopy()...)
-}
-
 func (f *fronted) providerFor(m Front) *Provider {
 	pid := m.getProviderID()
 	if pid == "" {

fronted_test.go

@@ -155,7 +155,7 @@ func doTestDomainFronting(t *testing.T, cacheFile string, expectedMasqueradesAtE
 	// Check the number of masquerades at the end, waiting until we get the right number
 	masqueradesAtEnd := 0
 	for i := 0; i < 1000; i++ {
-		masqueradesAtEnd = len(d.fronts)
+		masqueradesAtEnd = len(d.fronts.fronts)
 		if masqueradesAtEnd == expectedMasqueradesAtEnd {
 			break
 		}
@@ -761,9 +761,10 @@ func TestFindWorkingMasquerades(t *testing.T) {
 			}
 			f.providers = make(map[string]*Provider)
 			f.providers["testProviderId"] = NewProvider(nil, "", nil, nil, nil, nil, "")
-			f.fronts = make(sortedFronts, len(tt.masquerades))
+			//f.fronts = make(sortedFronts, len(tt.masquerades))
+			f.fronts = newThreadSafeFronts(len(tt.masquerades))
 			for i, m := range tt.masquerades {
-				f.fronts[i] = m
+				f.fronts.fronts[i] = m
 			}
 
 			f.tryAllFronts()