github.com/Azareal/Gosora@v0.0.0-20210729070923-553e66b59003/common/counters/referrers.go (about)

     1  package counters
     2  
     3  import (
     4  	"database/sql"
     5  	"sync"
     6  	"sync/atomic"
     7  
     8  	c "github.com/Azareal/Gosora/common"
     9  	qgen "github.com/Azareal/Gosora/query_gen"
    10  	"github.com/pkg/errors"
    11  )
    12  
    13  var ReferrerTracker *DefaultReferrerTracker
    14  
    15  // Add ReferrerItems here after they've had zero views for a while
    16  var referrersToDelete = make(map[string]*ReferrerItem)
    17  
    18  type ReferrerItem struct {
    19  	Count int64
    20  }
    21  
    22  // ? We'll track referrer domains here rather than the exact URL they arrived from for now, we'll think about expanding later
    23  // ? Referrers are fluid and ever-changing so we have to use string keys rather than 'enum' ints
    24  type DefaultReferrerTracker struct {
    25  	odd      map[string]*ReferrerItem
    26  	even     map[string]*ReferrerItem
    27  	oddLock  sync.RWMutex
    28  	evenLock sync.RWMutex
    29  
    30  	insert *sql.Stmt
    31  }
    32  
    33  func NewDefaultReferrerTracker() (*DefaultReferrerTracker, error) {
    34  	acc := qgen.NewAcc()
    35  	refTracker := &DefaultReferrerTracker{
    36  		odd:    make(map[string]*ReferrerItem),
    37  		even:   make(map[string]*ReferrerItem),
    38  		insert: acc.Insert("viewchunks_referrers").Columns("count,createdAt,domain").Fields("?,UTC_TIMESTAMP(),?").Prepare(), // TODO: Do something more efficient than doing a query for each referrer
    39  	}
    40  	c.Tasks.FifteenMin.Add(refTracker.Tick)
    41  	//c.Tasks.Sec.Add(refTracker.Tick)
    42  	c.Tasks.Shutdown.Add(refTracker.Tick)
    43  	return refTracker, acc.FirstError()
    44  }
    45  
    46  // TODO: Move this and the other view tickers out of the main task loop to avoid blocking other tasks?
    47  func (ref *DefaultReferrerTracker) Tick() (err error) {
    48  	for referrer, counter := range referrersToDelete {
    49  		// Handle views which squeezed through the gaps at the last moment
    50  		count := counter.Count
    51  		if count != 0 {
    52  			err := ref.insertChunk(referrer, count) // TODO: Bulk insert for speed?
    53  			if err != nil {
    54  				return errors.Wrap(errors.WithStack(err), "ref counter")
    55  			}
    56  		}
    57  		delete(referrersToDelete, referrer)
    58  	}
    59  
    60  	//  Run the queries and schedule zero view refs for deletion from memory
    61  	refLoop := func(l *sync.RWMutex, m map[string]*ReferrerItem) error {
    62  		l.Lock()
    63  		defer l.Unlock()
    64  		for referrer, counter := range m {
    65  			if counter.Count == 0 {
    66  				referrersToDelete[referrer] = counter
    67  				delete(m, referrer)
    68  			}
    69  			count := atomic.SwapInt64(&counter.Count, 0)
    70  			err := ref.insertChunk(referrer, count) // TODO: Bulk insert for speed?
    71  			if err != nil {
    72  				return errors.Wrap(errors.WithStack(err), "ref counter")
    73  			}
    74  		}
    75  		return nil
    76  	}
    77  	err = refLoop(&ref.oddLock, ref.odd)
    78  	if err != nil {
    79  		return err
    80  	}
    81  	return refLoop(&ref.evenLock, ref.even)
    82  }
    83  
    84  func (ref *DefaultReferrerTracker) insertChunk(referrer string, count int64) error {
    85  	if count == 0 {
    86  		return nil
    87  	}
    88  	c.DebugDetailf("Inserting a vchunk with a count of %d for ref %s", count, referrer)
    89  	_, err := ref.insert.Exec(count, referrer)
    90  	return err
    91  }
    92  
    93  func (ref *DefaultReferrerTracker) Bump(referrer string) {
    94  	if referrer == "" {
    95  		return
    96  	}
    97  	var refItem *ReferrerItem
    98  
    99  	// Slightly crude and rudimentary, but it should give a basic degree of sharding
   100  	if referrer[0]%2 == 0 {
   101  		ref.evenLock.RLock()
   102  		refItem = ref.even[referrer]
   103  		ref.evenLock.RUnlock()
   104  		if refItem != nil {
   105  			atomic.AddInt64(&refItem.Count, 1)
   106  		} else {
   107  			ref.evenLock.Lock()
   108  			ref.even[referrer] = &ReferrerItem{Count: 1}
   109  			ref.evenLock.Unlock()
   110  		}
   111  	} else {
   112  		ref.oddLock.RLock()
   113  		refItem = ref.odd[referrer]
   114  		ref.oddLock.RUnlock()
   115  		if refItem != nil {
   116  			atomic.AddInt64(&refItem.Count, 1)
   117  		} else {
   118  			ref.oddLock.Lock()
   119  			ref.odd[referrer] = &ReferrerItem{Count: 1}
   120  			ref.oddLock.Unlock()
   121  		}
   122  	}
   123  }