6
6
"io"
7
7
"math/rand"
8
8
"net/http"
9
+ "net/textproto"
9
10
"path"
10
11
"strings"
11
12
"sync"
@@ -39,10 +40,12 @@ type Distributor struct {
39
40
alertmanagerClientsPool ClientsPool
40
41
41
42
logger log.Logger
43
+
44
+ targetHeaders []string
42
45
}
43
46
44
47
// NewDistributor constructs a new Distributor
45
- func NewDistributor (cfg ClientConfig , maxRecvMsgSize int64 , alertmanagersRing * ring.Ring , alertmanagerClientsPool ClientsPool , logger log.Logger , reg prometheus.Registerer ) (d * Distributor , err error ) {
48
+ func NewDistributor (cfg ClientConfig , maxRecvMsgSize int64 , alertmanagersRing * ring.Ring , alertmanagerClientsPool ClientsPool , targetHeaders [] string , logger log.Logger , reg prometheus.Registerer ) (d * Distributor , err error ) {
46
49
if alertmanagerClientsPool == nil {
47
50
alertmanagerClientsPool = newAlertmanagerClientsPool (client .NewRingServiceDiscovery (alertmanagersRing ), cfg , logger , reg )
48
51
}
@@ -53,6 +56,7 @@ func NewDistributor(cfg ClientConfig, maxRecvMsgSize int64, alertmanagersRing *r
53
56
maxRecvMsgSize : maxRecvMsgSize ,
54
57
alertmanagerRing : alertmanagersRing ,
55
58
alertmanagerClientsPool : alertmanagerClientsPool ,
59
+ targetHeaders : targetHeaders ,
56
60
}
57
61
58
62
d .Service = services .NewBasicService (nil , d .running , nil )
@@ -170,6 +174,7 @@ func (d *Distributor) doQuorum(userID string, w http.ResponseWriter, r *http.Req
170
174
var responses []* httpgrpc.HTTPResponse
171
175
var responsesMtx sync.Mutex
172
176
grpcHeaders := httpToHttpgrpcHeaders (r .Header )
177
+
173
178
err = ring .DoBatch (r .Context (), RingOp , d .alertmanagerRing , []uint32 {shardByUser (userID )}, func (am ring.InstanceDesc , _ []int ) error {
174
179
// Use a background context to make sure all alertmanagers get the request even if we return early.
175
180
localCtx := opentracing .ContextWithSpan (user .InjectOrgID (context .Background (), userID ), opentracing .SpanFromContext (r .Context ()))
@@ -305,6 +310,23 @@ func (d *Distributor) doRequest(ctx context.Context, am ring.InstanceDesc, req *
305
310
return nil , errors .Wrapf (err , "failed to get alertmanager client from pool (alertmanager address: %s)" , am .Addr )
306
311
}
307
312
313
+ headers := make (map [string ]string , 0 )
314
+ for _ , h := range req .Headers {
315
+ headers [h .Key ] = h .Values [0 ]
316
+ }
317
+
318
+ headerMap := make (map [string ]string , 0 )
319
+ // Remove non-existent header.
320
+ for _ , header := range d .targetHeaders {
321
+ if v , ok := headers [textproto .CanonicalMIMEHeaderKey (header )]; ok {
322
+ headerMap [header ] = v
323
+ }
324
+ }
325
+
326
+ if len (headerMap ) > 0 {
327
+ ctx = util_log .ContextWithHeaderMap (ctx , headerMap )
328
+ }
329
+
308
330
return amClient .HandleRequest (ctx , req )
309
331
}
310
332
0 commit comments