121 lines
2.5 KiB
Go
121 lines
2.5 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"net/http/httputil"
|
|
"net/url"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/caarlos0/env/v11"
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/prometheus/client_golang/prometheus/promauto"
|
|
"github.com/prometheus/client_golang/prometheus/promhttp"
|
|
)
|
|
|
|
var cfg config
|
|
|
|
type config struct {
|
|
BaseURL string `env:"BASE_URL"`
|
|
Port int `env:"PORT"`
|
|
}
|
|
|
|
func createProxy(target *url.URL) func(http.ResponseWriter, *http.Request) {
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
|
r.URL.Path = strings.TrimPrefix(r.URL.Path, "/proxy")
|
|
r.Host = target.Host
|
|
r.URL.Scheme = target.Scheme
|
|
r.URL.Host = target.Host
|
|
|
|
director := func(req *http.Request) {
|
|
req.Header.Set("X-Forwarded-For", r.RemoteAddr)
|
|
req.URL = r.URL
|
|
req.Host = target.Host
|
|
req.URL.Scheme = target.Scheme
|
|
req.URL.Host = target.Host
|
|
}
|
|
|
|
modifyResponse := func(response *http.Response) error {
|
|
pr, pw := io.Pipe()
|
|
body := response.Body
|
|
response.Body = pr
|
|
|
|
go func() {
|
|
defer pw.Close()
|
|
|
|
reader := bufio.NewReader(body)
|
|
for {
|
|
line, err := reader.ReadBytes('\n')
|
|
if err != nil {
|
|
if err == io.EOF {
|
|
handleJsonLine([]byte(string(line)))
|
|
pw.Write(line)
|
|
break
|
|
}
|
|
return
|
|
}
|
|
handleJsonLine(line)
|
|
pw.Write(line)
|
|
}
|
|
|
|
}()
|
|
|
|
return nil
|
|
}
|
|
|
|
proxy := httputil.NewSingleHostReverseProxy(target)
|
|
proxy.Director = director
|
|
if r.URL.Path == "/api/generate" {
|
|
proxy.ModifyResponse = modifyResponse
|
|
}
|
|
|
|
proxy.ServeHTTP(w, r)
|
|
}
|
|
}
|
|
|
|
func handleJsonLine(line []byte) {
|
|
var jsonData map[string]interface{}
|
|
err := json.Unmarshal([]byte(line), &jsonData)
|
|
if err != nil {
|
|
fmt.Println("Error parsing JSON:", err)
|
|
return
|
|
}
|
|
|
|
if jsonData["done"].(bool) {
|
|
duration := jsonData["eval_duration"].(float64)
|
|
fmt.Printf("Duration: %.2f seconds\n", duration/1000000000.0)
|
|
opsProcessed.Inc()
|
|
}
|
|
}
|
|
|
|
var opsProcessed = promauto.NewCounter(prometheus.CounterOpts{
|
|
Name: "total_requests",
|
|
Help: "The total number of processed events",
|
|
})
|
|
|
|
func main() {
|
|
err := env.Parse(&cfg)
|
|
if err != nil {
|
|
log.Fatalf("Error parsing environment variables: %v", err)
|
|
}
|
|
|
|
targetURL, err := url.Parse(cfg.BaseURL)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
http.HandleFunc("/proxy/", createProxy(targetURL))
|
|
http.Handle("/metrics", promhttp.Handler())
|
|
|
|
log.Printf("Starting proxy server on :%s", strconv.Itoa(cfg.Port))
|
|
err = http.ListenAndServe(fmt.Sprintf(":%d", cfg.Port), nil)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
}
|