bouncer/internal/bench/proxy_test.go

package proxy_test

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"os"
	"path/filepath"
	"runtime/pprof"
	"strings"
	"testing"
	"time"

	"forge.cadoles.com/Cadoles/go-proxy"
	"forge.cadoles.com/cadoles/bouncer/internal/cache/memory"
	"forge.cadoles.com/cadoles/bouncer/internal/cache/ttl"
	"forge.cadoles.com/cadoles/bouncer/internal/config"
	"forge.cadoles.com/cadoles/bouncer/internal/proxy/director"
	"forge.cadoles.com/cadoles/bouncer/internal/store"
	redisStore "forge.cadoles.com/cadoles/bouncer/internal/store/redis"
	"github.com/pkg/errors"
	"github.com/redis/go-redis/v9"
	"gitlab.com/wpetit/goweb/logger"
	"gopkg.in/yaml.v3"

	"forge.cadoles.com/cadoles/bouncer/internal/setup"
)

func BenchmarkProxies(b *testing.B) {
	proxyFiles, err := filepath.Glob("testdata/proxies/*.yml")
	if err != nil {
		b.Fatalf("%+v", errors.WithStack(err))
	}

	for _, f := range proxyFiles {
		name := strings.TrimSuffix(filepath.Base(f), filepath.Ext(f))

		b.Run(name, func(b *testing.B) {
			heap, err := os.Create(filepath.Join("testdata", "proxies", name+"_heap.prof"))
			if err != nil {
				b.Fatalf("%+v", errors.Wrapf(err, "could not create heap profile"))
			}

			defer func() {
				defer heap.Close()

				if err := pprof.WriteHeapProfile(heap); err != nil {
					b.Fatalf("%+v", errors.WithStack(err))
				}
			}()

			conf, err := loadProxyBenchConfig(f)
			if err != nil {
				b.Fatalf("%+v", errors.Wrapf(err, "could notre load bench config"))
			}

			proxy, backend, err := createProxy(name, conf, b.Logf)
			if err != nil {
				b.Fatalf("%+v", errors.Wrapf(err, "could not create proxy"))
			}

			defer proxy.Close()

			if backend != nil {
				defer backend.Close()
			}

			client := proxy.Client()

			proxyURL, err := url.Parse(proxy.URL)
			if err != nil {
				b.Fatalf("%+v", errors.Wrapf(err, "could not parse proxy url"))
			}

			if conf.Fetch.URL.Path != "" {
				proxyURL.Path = conf.Fetch.URL.Path
			}

			if conf.Fetch.URL.RawQuery != "" {
				proxyURL.RawQuery = conf.Fetch.URL.RawQuery
			}

			if conf.Fetch.URL.User.Username != "" || conf.Fetch.URL.User.Password != "" {
				proxyURL.User = url.UserPassword(conf.Fetch.URL.User.Username, conf.Fetch.URL.User.Password)
			}

			rawProxyURL := proxyURL.String()

			b.Logf("fetching url '%s'", rawProxyURL)

			profile, err := os.Create(filepath.Join("testdata", "proxies", name+"_cpu.prof"))
			if err != nil {
				b.Fatalf("%+v", errors.Wrapf(err, "could not create cpu profile"))
			}

			defer profile.Close()

			if err := pprof.StartCPUProfile(profile); err != nil {
				b.Fatalf("%+v", errors.WithStack(err))
			}

			defer pprof.StopCPUProfile()

			b.ResetTimer()

			for i := 0; i < b.N; i++ {
				res, err := client.Get(rawProxyURL)
				if err != nil {
					b.Errorf("could not fetch proxy url: %+v", errors.WithStack(err))
				}

				body, err := io.ReadAll(res.Body)
				if err != nil {
					b.Errorf("could not read response body: %+v", errors.WithStack(err))
				}

				b.Logf("%s \n %v", res.Status, string(body))

				if err := res.Body.Close(); err != nil {
					b.Errorf("could not close response body: %+v", errors.WithStack(err))
				}
			}
		})
	}
}

type proxyBenchConfig struct {
	Proxy config.BootstrapProxyConfig `yaml:"proxy"`
	Fetch fetchBenchConfig            `yaml:"fetch"`
}

type fetchBenchConfig struct {
	URL fetchURLBenchConfig `yaml:"url"`
}

type fetchURLBenchConfig struct {
	Path     string                  `yaml:"path"`
	RawQuery string                  `yaml:"rawQuery"`
	User     fetchURLUserBenchConfig `yaml:"user"`
}

type fetchURLUserBenchConfig struct {
	Username string `yaml:"username"`
	Password string `yaml:"password"`
}

func loadProxyBenchConfig(filename string) (*proxyBenchConfig, error) {
	data, err := os.ReadFile(filename)
	if err != nil {
		return nil, errors.Wrapf(err, "could not read file '%s'", filename)
	}

	conf := proxyBenchConfig{}

	if err := yaml.Unmarshal(data, &conf); err != nil {
		return nil, errors.Wrapf(err, "could not unmarshal config")
	}

	return &conf, nil
}

func createProxy(name string, conf *proxyBenchConfig, logf func(format string, a ...any)) (*httptest.Server, *httptest.Server, error) {
	redisEndpoint := os.Getenv("BOUNCER_BENCH_REDIS_ADDR")
	if redisEndpoint == "" {
		redisEndpoint = "127.0.0.1:6379"
	}

	client := redis.NewUniversalClient(&redis.UniversalOptions{
		Addrs: []string{redisEndpoint},
	})

	proxyRepository := redisStore.NewProxyRepository(client, redisStore.DefaultTxMaxAttempts, redisStore.DefaultTxBaseDelay)
	layerRepository := redisStore.NewLayerRepository(client, redisStore.DefaultTxMaxAttempts, redisStore.DefaultTxBaseDelay)

	var backend *httptest.Server

	if conf.Proxy.To == "" {
		backend = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
			if _, err := w.Write([]byte("Hello, world.")); err != nil {
				logf("[ERROR] %+v", errors.WithStack(err))
			}
		}))

		if err := waitFor(backend.URL, 5*time.Second); err != nil {
			return nil, nil, errors.WithStack(err)
		}

		logf("started backend '%s'", backend.URL)
	}

	ctx := context.Background()

	proxyName := store.ProxyName("bench-" + name)

	proxies, err := proxyRepository.QueryProxy(ctx)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}

	// Cleanup existing proxies
	for _, p := range proxies {
		if err := proxyRepository.DeleteProxy(ctx, p.Name); err != nil {
			return nil, nil, errors.WithStack(err)
		}
	}

	logf("creating proxy '%s'", proxyName)

	to := string(conf.Proxy.To)
	if to == "" {
		to = backend.URL
	}

	if _, err := proxyRepository.CreateProxy(ctx, proxyName, to, conf.Proxy.From...); err != nil {
		return nil, nil, errors.WithStack(err)
	}

	if _, err := proxyRepository.UpdateProxy(ctx, proxyName, store.WithProxyUpdateEnabled(true)); err != nil {
		return nil, nil, errors.WithStack(err)
	}

	for layerName, layerConf := range conf.Proxy.Layers {
		if err := layerRepository.DeleteLayer(ctx, proxyName, store.LayerName(layerName)); err != nil {
			return nil, nil, errors.WithStack(err)
		}

		_, err := layerRepository.CreateLayer(ctx, proxyName, store.LayerName(layerName), store.LayerType(layerConf.Type), layerConf.Options.Data)
		if err != nil {
			return nil, nil, errors.WithStack(err)
		}

		_, err = layerRepository.UpdateLayer(ctx, proxyName, store.LayerName(layerName), store.WithLayerUpdateEnabled(bool(layerConf.Enabled)))
		if err != nil {
			return nil, nil, errors.WithStack(err)
		}

	}

	appConf := config.NewDefault()
	appConf.Logger.Level = config.InterpolatedInt(logger.LevelError)
	appConf.Layers.Authn.TemplateDir = "../../layers/authn/templates"
	appConf.Layers.Queue.TemplateDir = "../../layers/queue/templates"

	layers, err := setup.GetLayers(context.Background(), appConf)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}

	director := director.New(
		proxyRepository, layerRepository,
		director.WithLayerCache(
			ttl.NewCache(
				memory.NewCache[string, []*store.Layer](),
				memory.NewCache[string, time.Time](),
				30*time.Second,
			),
		),
		director.WithProxyCache(
			ttl.NewCache(
				memory.NewCache[string, []*store.Proxy](),
				memory.NewCache[string, time.Time](),
				30*time.Second,
			),
		),
		director.WithLayers(layers...),
	)

	directorMiddleware := director.Middleware()

	handler := proxy.New(
		proxy.WithRequestTransformers(
			director.RequestTransformer(),
		),
		proxy.WithResponseTransformers(
			director.ResponseTransformer(),
		),
		proxy.WithReverseProxyFactory(func(ctx context.Context, target *url.URL) *httputil.ReverseProxy {
			reverse := httputil.NewSingleHostReverseProxy(target)
			reverse.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
				logf("[ERROR] %s", errors.WithStack(err))
			}
			return reverse
		}),
	)

	server := httptest.NewServer(directorMiddleware(handler))

	return server, backend, nil
}

func waitFor(url string, ttl time.Duration) error {
	var lastErr error
	timeout := time.After(ttl)
	for {
		select {
		case <-timeout:
			if lastErr != nil {
				return lastErr
			}

			return errors.New("wait timed out")
		default:
			res, err := http.Get(url)
			if err != nil {
				lastErr = errors.WithStack(err)
				continue
			}

			if res.StatusCode >= 200 && res.StatusCode < 400 {
				return nil
			}
		}
	}
}