irc-client/throttle.go

package ircclient

import (
	"log"
	"time"
)

// Track when the client last sent something
type FloodTrack struct {
	Pos     int // Index to store next item
	Size    int // Max number of Track items
	Timeout int // Timeout in seconds
	Track   []time.Time
}

// Initialize flood tracking
func (F *FloodTrack) Init(size int, timeout int) {
	F.Size = size
	F.Timeout = timeout
	F.Track = make([]time.Time, size)
}

// Are we full?
//
//	Expire records before checking.
func (F *FloodTrack) Full() bool {
	if F.Pos > 0 {
		F.Expire()
	}
	return F.Pos == F.Size
}

// Expire Track records older then timeout
func (F *FloodTrack) Expire() {
	var idx int

ReCheck:
	for idx = 0; idx < F.Pos; idx++ {
		// log.Println(idx, time.Since(F.Track[idx]).Seconds())
		if time.Since(F.Track[idx]).Seconds() > float64(F.Timeout) {
			// Remove this from the list
			F.Pos--
			for pos := idx; pos < F.Pos; pos++ {
				F.Track[pos] = F.Track[pos+1]
			}
			goto ReCheck
		}
	}
}

// Save Now() into the tracker.
func (F *FloodTrack) Save() {
	F.Track[F.Pos] = time.Now()
	F.Pos++
}

type ThrottleBuffer struct {
	buffer     map[string][]string // Map of targets and strings to send
	targets    []string            // List of targets
	last       int                 // Last index set
	Life_sucks bool                // Is throttle active?
}

func (T *ThrottleBuffer) Init() {
	T.buffer = make(map[string][]string, 0)
	T.targets = make([]string, 0)
	T.last = 0
}

// Pop next available from buffer.
//
//	If empty, remove from map and targets list.  Adjust last if needed.
func (T *ThrottleBuffer) pop() string {
	// Get next target
	var t string = T.targets[T.last]

	T.last++
	if T.last == len(T.targets) {
		// We're past the end, start over at the beginning.
		T.last = 0
	}

	var msg string = T.buffer[t][0]

	if len(T.buffer[t]) == 1 {
		// This is the last entry for this target.
		delete(T.buffer, t)

		if len(T.targets) == 1 {
			// This is the last entry
			T.targets = make([]string, 0)
			T.Life_sucks = false
			log.Println("Flood control off.")
		} else {
			// Remove t from targets
			for x := 0; x < len(T.targets); x++ {
				if T.targets[x] == t {
					T.targets = append(T.targets[:x], T.targets[x+1:]...)
					if x <= T.last {
						T.last--
					}
					break
				}

			}
		}
	} else {
		// Delete the first entry for t
		T.buffer[t] = append(T.buffer[t][:0], T.buffer[t][1:]...)
	}

	return msg
}

// Push item into the buffer, begin throttling.
func (T *ThrottleBuffer) push(To string, Output string) {
	if len(T.targets) == 0 {
		T.Life_sucks = true
		T.last = 0
		log.Println("Flood control enabled.")
	}

	_, has := T.buffer[To]
	if !has {
		T.buffer[To] = make([]string, 0)
		T.targets = append(T.targets, To)
	}
	T.buffer[To] = append(T.buffer[To], Output)
}

// Delete target from buffer, if exists.
//
//	If we're kicked from the channel we're sending to, or if the
//	nick we're sending to quits -- remove those from write queue.
func (T *ThrottleBuffer) delete(To string) {
	_, has := T.buffer[To]
	if has {
		// Yes, the buffer has message(s) for To
		delete(T.buffer, To)
		if len(T.targets) == 1 {
			// Last entry
			T.targets = make([]string, 0)
			T.Life_sucks = false
			log.Println("Flood control off.")
		} else {
			// Remove To from targets
			for x := 0; x < len(T.targets); x++ {
				if T.targets[x] == To {
					T.targets = append(T.targets[:x], T.targets[x+1:]...)
					// Don't decrement if already at first item
					if (x <= T.last) && (T.last != 0) {
						T.last--
					}
					break
				}
			}
		}
	}
}