bouncer/internal/proxy/proxy_test.go

package proxy_test

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"os"
	"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/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"
)

func BenchmarkProxy(b *testing.B) {
	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)
	layerRepository := redisStore.NewLayerRepository(client)

	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 {
			b.Logf("[ERROR] %+v", errors.WithStack(err))
		}
	}))
	defer backend.Close()

	if err := waitFor(backend.URL, 5*time.Second); err != nil {
		b.Fatalf("[FATAL] %+v", errors.WithStack(err))
	}

	b.Logf("started backend '%s'", backend.URL)

	ctx := context.Background()

	proxyName := store.ProxyName(b.Name())

	b.Logf("creating proxy '%s'", proxyName)

	if err := proxyRepository.DeleteProxy(ctx, proxyName); err != nil {
		b.Fatalf("[FATAL] %+v", errors.WithStack(err))
	}

	if _, err := proxyRepository.CreateProxy(ctx, proxyName, backend.URL, "*"); err != nil {
		b.Fatalf("[FATAL] %+v", errors.WithStack(err))
	}

	if _, err := proxyRepository.UpdateProxy(ctx, proxyName, store.WithProxyUpdateEnabled(true)); err != nil {
		b.Fatalf("[FATAL] %+v", 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,
			),
		),
	)

	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) {
				b.Logf("[ERROR] %s", errors.WithStack(err))
			}
			return reverse
		}),
	)

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

	b.Logf("started proxy '%s'", server.URL)

	httpClient := server.Client()

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		res, err := httpClient.Get(server.URL)
		if err != nil {
			b.Errorf("could not fetch server 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 - %v", res.Status, string(body))

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

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