overloader

README

overloader

A plugin to protect teleport from overload.

Test

package overloader

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"

	tp "github.com/henrylee2cn/teleport/v6"
	"github.com/stretchr/testify/assert"
)

type Home struct {
	tp.CallCtx
}

func (h *Home) Test(arg *map[string]string) (map[string]interface{}, *tp.Status) {
	return map[string]interface{}{
		"arg": *arg,
	}, nil
}

func TestPlugin(t *testing.T) {
	ol := New(LimitConfig{
		MaxConn:     1,
		QPSInterval: 100 * time.Millisecond,
		MaxTotalQPS: 2,
		MaxHandlerQPS: []HandlerLimit{
			{ServiceMethod: "/home/test", MaxQPS: 1},
		},
	})
	// Server
	srv := tp.NewPeer(
		tp.PeerConfig{ListenPort: 9090, CountTime: true},
		ol,
	)
	srv.RouteCall(new(Home))
	go srv.ListenAndServe()
	time.Sleep(1e9)

	// Client
	cli := tp.NewPeer(
		tp.PeerConfig{CountTime: true},
	)
	var testClient = func(connNum, totalQPS int) (olConnCount, olQPSCount int64) {
		var connGW sync.WaitGroup
		connGW.Add(connNum)
		defer connGW.Wait()
		for index := 0; index < connNum; index++ {
			go func() {
				defer connGW.Done()
				sess, stat := cli.Dial(":9090")
				if !stat.OK() {
					t.Fatal(stat)
				}
				defer sess.Close()
				time.Sleep(time.Millisecond)
				if !sess.Health() {
					atomic.AddInt64(&olConnCount, 1)
					t.Logf("connNum:%d, totalQPS:%d, dial: Connection Closed", connNum, totalQPS)
					return
				}
				var qpsGW sync.WaitGroup
				qpsGW.Add(totalQPS)
				defer qpsGW.Wait()
				for i := 0; i < totalQPS; i++ {
					go func() {
						defer qpsGW.Done()
						stat = sess.Call("/home/test", nil, nil).Status()
						if !stat.OK() {
							atomic.AddInt64(&olQPSCount, 1)
							t.Logf("connNum:%d, totalQPS:%d, call:%s", connNum, totalQPS, stat)
						}
					}()
				}
			}()
		}
		return
	}

	{
		olConnCount, olQPSCount := testClient(1, 1)
		assert.Equal(t, int64(0), olConnCount)
		assert.Equal(t, int64(0), olQPSCount)
	}
	{
		time.Sleep(time.Second)
		assert.Equal(t, 0, srv.CountSession())
		olConnCount, olQPSCount := testClient(2, 1)
		assert.Equal(t, int64(1), olConnCount)
		assert.Equal(t, int64(0), olQPSCount)
	}
	{
		time.Sleep(time.Second)
		assert.Equal(t, 0, srv.CountSession())
		olConnCount, olQPSCount := testClient(1, 2)
		assert.Equal(t, int64(0), olConnCount)
		assert.Equal(t, int64(1), olQPSCount)
	}
	{
		ol.Update(LimitConfig{
			MaxConn:     100,
			QPSInterval: time.Second,
			MaxTotalQPS: 2000,
		})
		time.Sleep(time.Second) // Wait for one cycle to take effect
		assert.Equal(t, 0, srv.CountSession())
		olConnCount, olQPSCount := testClient(10, 200)
		assert.Equal(t, int64(0), olConnCount)
		assert.Equal(t, int64(0), olQPSCount)
	}
}

test command:

go test -v -run=TestPlugin

Documentation

Overview

Package overloader is a plugin to protect teleport from overload.

Index

• type HandlerLimit
• type LimitConfig
• type Overloader
  • func New(initLimitConfig LimitConfig) *Overloader
  • func (o *Overloader) LimitConfig() LimitConfig
  • func (o *Overloader) Name() string
  • func (o *Overloader) PostAccept(_ tp.PreSession) *tp.Status
  • func (o *Overloader) PostDial(sess tp.PreSession, isRedial bool) *tp.Status
  • func (o *Overloader) PostDisconnect(_ tp.BaseSession) *tp.Status
  • func (o *Overloader) PostReadCallHeader(ctx tp.ReadCtx) *tp.Status
  • func (o *Overloader) PostReadPushHeader(ctx tp.ReadCtx) *tp.Status
  • func (o *Overloader) Update(newLimitConfig LimitConfig)

Types

type HandlerLimit

type HandlerLimit struct {
	ServiceMethod string
	MaxQPS        int32
}

HandlerLimit handler QPS overload limitation condition

type LimitConfig

type LimitConfig struct {
	MaxConn       int32
	QPSInterval   time.Duration
	MaxTotalQPS   int32
	MaxHandlerQPS []HandlerLimit
}

LimitConfig overload limitation condition

type Overloader

type Overloader struct {
	// contains filtered or unexported fields
}

Overloader plug-in to protect teleport from overload

func New

func New(initLimitConfig LimitConfig) *Overloader

New creates a plug-in to protect teleport from overload.

func (*Overloader) LimitConfig

func (o *Overloader) LimitConfig() LimitConfig

LimitConfig returns the overload limitation condition.

func (*Overloader) Name

func (o *Overloader) Name() string

Name returns the plugin name.

func (*Overloader) PostAccept

func (o *Overloader) PostAccept(_ tp.PreSession) *tp.Status

PostAccept checks connection overload. If overload, print error log and close the connection.

func (*Overloader) PostDial

func (o *Overloader) PostDial(sess tp.PreSession, isRedial bool) *tp.Status

PostDial checks connection overload. If overload, print error log and close the connection.

func (*Overloader) PostDisconnect

func (o *Overloader) PostDisconnect(_ tp.BaseSession) *tp.Status

PostDisconnect releases connection count.

func (*Overloader) PostReadCallHeader

func (o *Overloader) PostReadCallHeader(ctx tp.ReadCtx) *tp.Status

PostReadCallHeader checks PULL QPS overload. If overload, print error log and reply error.

func (*Overloader) PostReadPushHeader

func (o *Overloader) PostReadPushHeader(ctx tp.ReadCtx) *tp.Status

PostReadPushHeader checks PUSH QPS overload. If overload, print warning log.

func (*Overloader) Update

func (o *Overloader) Update(newLimitConfig LimitConfig)

Update updates the overload limitation condition.