irc-client/throttle.go

package ircclient

import (
	"log"
	"time"
)

type FloodTrack struct {
	Pos     int
	Size    int
	Timeout int
	Track   []time.Time
}

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

func (F *FloodTrack) Full() bool {
	if F.Pos > 0 {
		F.Expire()
	}
	return F.Pos == F.Size
}

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
		}
	}
}

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

type ThrottleBuffer struct {
	buffer     map[string][]string
	targets    []string
	last       int
	Life_sucks bool
}

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
}

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)
}

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
				}
			}
		}
	}
}