17612909769

Committed 10 Sep 2025 11:55AM UTC coverage: 69.676% (+1.4%) from 68.32%

Build # 17612909769

Build Type

Pull #94

github

Committed by

MrWong99

Commit Message

feat: added loki streams processing at /loki/push

Signed-off-by: Lukas Schmidt <luk.schm@web.de>

Pull Request Pull Request #94: feat: added loki streams processing at POST /loki/push

Coverage Stats

266 of 367 new or added lines in 5 files covered. (72.48%)

409 of 587 relevant lines covered (69.68%)

6.08 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

86.96

/processor.go

package main

import (
        "bytes"
        "encoding/base64"
        "fmt"
        "net"
        "sync"
        "sync/atomic"
        "time"

        "github.com/blind-oracle/go-common/logger"
        "github.com/google/uuid"
        fh "github.com/valyala/fasthttp"
)

type result struct {
        code     int
        body     []byte
        duration float64
        tenant   string
        err      error
}

type processor struct {
        cfg config

        srv *fh.Server
        cli *fh.Client

        shuttingDown uint32

        logger.Logger

        auth struct {
                egressHeader []byte
        }
}

func newProcessor(c config) *processor {
        p := &processor{
                cfg:    c,
                Logger: logger.NewSimpleLogger("proc"),
        }

        p.srv = &fh.Server{
                Name:    "cortex-tenant",
                Handler: p.handle,

                MaxRequestBodySize: 8 * 1024 * 1024,

                ReadTimeout:  c.Timeout,
                WriteTimeout: c.Timeout,
                IdleTimeout:  c.IdleTimeout,

                Concurrency: c.Concurrency,
        }

        p.cli = &fh.Client{
                Name:               "cortex-tenant",
                ReadTimeout:        c.Timeout,
                WriteTimeout:       c.Timeout,
                MaxConnWaitTimeout: 1 * time.Second,
                MaxConnsPerHost:    c.MaxConnsPerHost,
                DialDualStack:      c.EnableIPv6,
                MaxConnDuration:    c.MaxConnDuration,
        }

        if c.Auth.Egress.Username != "" {
                authString := []byte(fmt.Sprintf("%s:%s", c.Auth.Egress.Username, c.Auth.Egress.Password))
                p.auth.egressHeader = []byte("Basic " + base64.StdEncoding.EncodeToString(authString))
        }

        // For testing
        if c.pipeOut != nil {
                p.cli.Dial = func(a string) (net.Conn, error) {
                        return c.pipeOut.Dial()
                }
        }

        return p
}

func (p *processor) run() (err error) {
        var l net.Listener

        // For testing
        if p.cfg.pipeIn == nil {
                if l, err = net.Listen("tcp", p.cfg.Listen); err != nil {
                        return
                }
        } else {
                l = p.cfg.pipeIn
        }

        go p.srv.Serve(l)
        return
}

func (p *processor) handle(ctx *fh.RequestCtx) {
        if bytes.Equal(ctx.Path(), []byte("/alive")) {
                if atomic.LoadUint32(&p.shuttingDown) == 1 {
                        ctx.SetStatusCode(fh.StatusServiceUnavailable)
                }

                return
        }

        if !bytes.Equal(ctx.Request.Header.Method(), []byte("POST")) {
                ctx.Error("Expecting POST", fh.StatusBadRequest)
                return
        }

        if bytes.Equal(ctx.Path(), []byte("/push")) {
                p.handleMetrics(ctx)
                return
        }

        if bytes.Equal(ctx.Path(), []byte("/loki/push")) {
                p.handleLogs(ctx)
                return
        }

        ctx.SetStatusCode(fh.StatusNotFound)
}

func (p *processor) dispatch(target string, clientIP net.Addr, reqID uuid.UUID, tenantPrefix string, m map[string]func() ([]byte, error)) (res []result) {
        var wg sync.WaitGroup
        res = make([]result, len(m))

        i := 0
        for tenant, bodyFunc := range m {
                wg.Add(1)

                go func(idx int, tenant string, bodyFunc func() ([]byte, error)) {
                        defer wg.Done()

                        r := p.send(target, clientIP, reqID, tenant, bodyFunc)
                        res[idx] = r
                }(i, tenantPrefix+tenant, bodyFunc)

                i++
        }

        wg.Wait()
        return
}

func (p *processor) send(target string, clientIP net.Addr, reqID uuid.UUID, tenant string, bodyFunc func() ([]byte, error)) (r result) {
        start := time.Now()
        r.tenant = tenant

        req := fh.AcquireRequest()
        resp := fh.AcquireResponse()

        defer func() {
                fh.ReleaseRequest(req)
                fh.ReleaseResponse(resp)
        }()

        buf, err := bodyFunc()
        if err != nil {
                r.err = err
                return
        }

        p.fillRequestHeaders(clientIP, reqID, tenant, req)

        if p.auth.egressHeader != nil {
                req.Header.SetBytesV("Authorization", p.auth.egressHeader)
        }

        req.Header.SetMethod(fh.MethodPost)
        req.SetRequestURI(target)
        req.SetBody(buf)

        if err = p.cli.DoTimeout(req, resp, p.cfg.Timeout); err != nil {
                r.err = err
                return
        }

        r.code = resp.Header.StatusCode()
        r.body = make([]byte, len(resp.Body()))
        copy(r.body, resp.Body())
        r.duration = time.Since(start).Seconds() / 1000

        return
}

func (p *processor) fillRequestHeaders(
        clientIP net.Addr, reqID uuid.UUID, tenant string, req *fh.Request) {
        req.Header.Set("Content-Encoding", "snappy")
        req.Header.Set("Content-Type", "application/x-protobuf")
        req.Header.Set("X-Prometheus-Remote-Write-Version", "0.1.0")
        req.Header.Set("X-Cortex-Tenant-Client", clientIP.String())
        req.Header.Set("X-Cortex-Tenant-ReqID", reqID.String())
        req.Header.Set(p.cfg.Tenant.Header, tenant)
}

func (p *processor) close() (err error) {
        // Signal that we're shutting down
        atomic.StoreUint32(&p.shuttingDown, 1)
        // Let healthcheck detect that we're offline
        time.Sleep(p.cfg.TimeoutShutdown)
        // Shutdown
        return p.srv.Shutdown()
}

1	package main
2
3	import (
4	"bytes"
5	"encoding/base64"
6	"fmt"
7	"net"
8	"sync"
9	"sync/atomic"
10	"time"
11
12	"github.com/blind-oracle/go-common/logger"
13	"github.com/google/uuid"
14	fh "github.com/valyala/fasthttp"
15	)
16
17	type result struct {
18	code int
19	body []byte
20	duration float64
21	tenant string
22	err error
23	}
24
25	type processor struct {
26	cfg config
27
28	srv *fh.Server
29	cli *fh.Client
30
31	shuttingDown uint32
32
33	logger.Logger
34
35	auth struct {
36	egressHeader []byte
37	}
38	}
39
40	func newProcessor(c config) *processor {	11✔
41	p := &processor{	11✔
42	cfg: c,	11✔
43	Logger: logger.NewSimpleLogger("proc"),	11✔
44	}	11✔
45		11✔
46	p.srv = &fh.Server{	11✔
47	Name: "cortex-tenant",	11✔
48	Handler: p.handle,	11✔
49		11✔
50	MaxRequestBodySize: 8 * 1024 * 1024,	11✔
51		11✔
52	ReadTimeout: c.Timeout,	11✔
53	WriteTimeout: c.Timeout,	11✔
54	IdleTimeout: c.IdleTimeout,	11✔
55		11✔
56	Concurrency: c.Concurrency,	11✔
57	}	11✔
58		11✔
59	p.cli = &fh.Client{	11✔
60	Name: "cortex-tenant",	11✔
61	ReadTimeout: c.Timeout,	11✔
62	WriteTimeout: c.Timeout,	11✔
63	MaxConnWaitTimeout: 1 * time.Second,	11✔
64	MaxConnsPerHost: c.MaxConnsPerHost,	11✔
65	DialDualStack: c.EnableIPv6,	11✔
66	MaxConnDuration: c.MaxConnDuration,	11✔
67	}	11✔
68		11✔
69	if c.Auth.Egress.Username != "" {	11✔
70	authString := []byte(fmt.Sprintf("%s:%s", c.Auth.Egress.Username, c.Auth.Egress.Password))	×
71	p.auth.egressHeader = []byte("Basic " + base64.StdEncoding.EncodeToString(authString))	×
72	}	×
73
74	// For testing
75	if c.pipeOut != nil {	13✔
76	p.cli.Dial = func(a string) (net.Conn, error) {	6✔
77	return c.pipeOut.Dial()	4✔
78	}	4✔
79	}
80
81	return p	11✔
82	}
83
84	func (p *processor) run() (err error) {	2✔
85	var l net.Listener	2✔
86		2✔
87	// For testing	2✔
88	if p.cfg.pipeIn == nil {	2✔
89	if l, err = net.Listen("tcp", p.cfg.Listen); err != nil {	×
90	return	×
91	}	×
92	} else {	2✔
93	l = p.cfg.pipeIn	2✔
94	}	2✔
95
96	go p.srv.Serve(l)	2✔
97	return	2✔
98	}
99
100	func (p processor) handle(ctx fh.RequestCtx) {	15✔
101	if bytes.Equal(ctx.Path(), []byte("/alive")) {	17✔
102	if atomic.LoadUint32(&p.shuttingDown) == 1 {	4✔
103	ctx.SetStatusCode(fh.StatusServiceUnavailable)	2✔
104	}	2✔
105
106	return	2✔
107	}
108
109	if !bytes.Equal(ctx.Request.Header.Method(), []byte("POST")) {	13✔
110	ctx.Error("Expecting POST", fh.StatusBadRequest)	×
111	return	×
112	}	×
113
114	if bytes.Equal(ctx.Path(), []byte("/push")) {	20✔
115	p.handleMetrics(ctx)	7✔
116	return	7✔
117	}	7✔
118
119	if bytes.Equal(ctx.Path(), []byte("/loki/push")) {	12✔
120	p.handleLogs(ctx)	6✔
121	return	6✔
122	}	6✔
123
NEW 124	ctx.SetStatusCode(fh.StatusNotFound)	×
125	}
126
127	func (p *processor) dispatch(target string, clientIP net.Addr, reqID uuid.UUID, tenantPrefix string, m map[string]func() ([]byte, error)) (res []result) {	6✔
128	var wg sync.WaitGroup	6✔
129	res = make([]result, len(m))	6✔
130		6✔
131	i := 0	6✔
132	for tenant, bodyFunc := range m {	18✔
133	wg.Add(1)	12✔
134		12✔
135	go func(idx int, tenant string, bodyFunc func() ([]byte, error)) {	24✔
136	defer wg.Done()	12✔
137		12✔
138	r := p.send(target, clientIP, reqID, tenant, bodyFunc)	12✔
139	res[idx] = r	12✔
140	}(i, tenantPrefix+tenant, bodyFunc)	12✔
141
142	i++	12✔
143	}
144
145	wg.Wait()	6✔
146	return	6✔
147	}
148
149	func (p *processor) send(target string, clientIP net.Addr, reqID uuid.UUID, tenant string, bodyFunc func() ([]byte, error)) (r result) {	12✔
150	start := time.Now()	12✔
151	r.tenant = tenant	12✔
152		12✔
153	req := fh.AcquireRequest()	12✔
154	resp := fh.AcquireResponse()	12✔
155		12✔
156	defer func() {	24✔
157	fh.ReleaseRequest(req)	12✔
158	fh.ReleaseResponse(resp)	12✔
159	}()	12✔
160
161	buf, err := bodyFunc()	12✔
162	if err != nil {	12✔
163	r.err = err	×
164	return	×
165	}	×
166
167	p.fillRequestHeaders(clientIP, reqID, tenant, req)	12✔
168		12✔
169	if p.auth.egressHeader != nil {	12✔
170	req.Header.SetBytesV("Authorization", p.auth.egressHeader)	×
171	}	×
172
173	req.Header.SetMethod(fh.MethodPost)	12✔
174	req.SetRequestURI(target)	12✔
175	req.SetBody(buf)	12✔
176		12✔
177	if err = p.cli.DoTimeout(req, resp, p.cfg.Timeout); err != nil {	12✔
178	r.err = err	×
179	return	×
180	}	×
181
182	r.code = resp.Header.StatusCode()	12✔
183	r.body = make([]byte, len(resp.Body()))	12✔
184	copy(r.body, resp.Body())	12✔
185	r.duration = time.Since(start).Seconds() / 1000	12✔
186		12✔
187	return	12✔
188	}
189
190	func (p *processor) fillRequestHeaders(
191	clientIP net.Addr, reqID uuid.UUID, tenant string, req *fh.Request) {	13✔
192	req.Header.Set("Content-Encoding", "snappy")	13✔
193	req.Header.Set("Content-Type", "application/x-protobuf")	13✔
194	req.Header.Set("X-Prometheus-Remote-Write-Version", "0.1.0")	13✔
195	req.Header.Set("X-Cortex-Tenant-Client", clientIP.String())	13✔
196	req.Header.Set("X-Cortex-Tenant-ReqID", reqID.String())	13✔
197	req.Header.Set(p.cfg.Tenant.Header, tenant)	13✔
198	}	13✔
199
200	func (p *processor) close() (err error) {	2✔
201	// Signal that we're shutting down	2✔
202	atomic.StoreUint32(&p.shuttingDown, 1)	2✔
203	// Let healthcheck detect that we're offline	2✔
204	time.Sleep(p.cfg.TimeoutShutdown)	2✔
205	// Shutdown	2✔
206	return p.srv.Shutdown()	2✔
207	}	2✔

blind-oracle / cortex-tenant / 17612909769

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous