wings/server/console.go
2022-10-06 09:58:42 -06:00

83 lines
2.3 KiB
Go

package server
import (
"fmt"
"sync"
"time"
"github.com/mitchellh/colorstring"
"github.com/pterodactyl/wings/config"
"github.com/pterodactyl/wings/system"
)
// appName is a local cache variable to avoid having to make expensive copies of
// the configuration every time we need to send output along to the websocket for
// a server.
var appName string
var appNameSync sync.Once
// PublishConsoleOutputFromDaemon sends output to the server console formatted
// to appear correctly as being sent from Wings.
func (s *Server) PublishConsoleOutputFromDaemon(data string) {
appNameSync.Do(func() {
appName = config.Get().AppName
})
s.Events().Publish(
ConsoleOutputEvent,
colorstring.Color(fmt.Sprintf("[yellow][bold][%s Daemon]:[default] %s", appName, data)),
)
}
// Throttler returns the throttler instance for the server or creates a new one.
func (s *Server) Throttler() *ConsoleThrottle {
s.throttleOnce.Do(func() {
throttles := config.Get().Throttles
period := time.Duration(throttles.Period) * time.Millisecond
s.throttler = newConsoleThrottle(throttles.Lines, period)
s.throttler.strike = func() {
s.PublishConsoleOutputFromDaemon(fmt.Sprintf("Server is outputting console data too quickly -- throttling..."))
}
})
return s.throttler
}
type ConsoleThrottle struct {
limit *system.Rate
lock *system.Locker
strike func()
}
func newConsoleThrottle(lines uint64, period time.Duration) *ConsoleThrottle {
return &ConsoleThrottle{
limit: system.NewRate(lines, period),
lock: system.NewLocker(),
}
}
// Allow checks if the console is allowed to process more output data, or if too
// much has already been sent over the line. If there is too much output the
// strike callback function is triggered, but only if it has not already been
// triggered at this point in the process.
//
// If output is allowed, the lock on the throttler is released and the next time
// it is triggered the strike function will be re-executed.
func (ct *ConsoleThrottle) Allow() bool {
if !ct.limit.Try() {
if err := ct.lock.Acquire(); err == nil {
if ct.strike != nil {
ct.strike()
}
}
return false
}
ct.lock.Release()
return true
}
// Reset resets the console throttler internal rate limiter and overage counter.
func (ct *ConsoleThrottle) Reset() {
ct.limit.Reset()
}