mirror of
https://gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake.git
synced 2025-10-13 20:11:19 -04:00
Make the proxy to report the number of clients to the broker
So the assignment of proxies is based on the load. The number of clients is ronded down to 8. Existing proxies that doesn't report the number of clients will be distributed equaly to new proxies until they get 8 clients, that is okish as the existing proxies do have a maximum capacity of 10. Fixes #40048
This commit is contained in:
parent
74bdb85b30
commit
7a1857c42f
9 changed files with 165 additions and 77 deletions
|
@ -124,16 +124,18 @@ type ProxyPoll struct {
|
|||
id string
|
||||
proxyType string
|
||||
natType string
|
||||
clients int
|
||||
offerChannel chan *ClientOffer
|
||||
}
|
||||
|
||||
// Registers a Snowflake and waits for some Client to send an offer,
|
||||
// as part of the polling logic of the proxy handler.
|
||||
func (ctx *BrokerContext) RequestOffer(id string, proxyType string, natType string) *ClientOffer {
|
||||
func (ctx *BrokerContext) RequestOffer(id string, proxyType string, natType string, clients int) *ClientOffer {
|
||||
request := new(ProxyPoll)
|
||||
request.id = id
|
||||
request.proxyType = proxyType
|
||||
request.natType = natType
|
||||
request.clients = clients
|
||||
request.offerChannel = make(chan *ClientOffer)
|
||||
ctx.proxyPolls <- request
|
||||
// Block until an offer is available, or timeout which sends a nil offer.
|
||||
|
@ -146,7 +148,7 @@ func (ctx *BrokerContext) RequestOffer(id string, proxyType string, natType stri
|
|||
// client offer or nil on timeout / none are available.
|
||||
func (ctx *BrokerContext) Broker() {
|
||||
for request := range ctx.proxyPolls {
|
||||
snowflake := ctx.AddSnowflake(request.id, request.proxyType, request.natType)
|
||||
snowflake := ctx.AddSnowflake(request.id, request.proxyType, request.natType, request.clients)
|
||||
// Wait for a client to avail an offer to the snowflake.
|
||||
go func(request *ProxyPoll) {
|
||||
select {
|
||||
|
@ -174,10 +176,10 @@ func (ctx *BrokerContext) Broker() {
|
|||
// Create and add a Snowflake to the heap.
|
||||
// Required to keep track of proxies between providing them
|
||||
// with an offer and awaiting their second POST with an answer.
|
||||
func (ctx *BrokerContext) AddSnowflake(id string, proxyType string, natType string) *Snowflake {
|
||||
func (ctx *BrokerContext) AddSnowflake(id string, proxyType string, natType string, clients int) *Snowflake {
|
||||
snowflake := new(Snowflake)
|
||||
snowflake.id = id
|
||||
snowflake.clients = 0
|
||||
snowflake.clients = clients
|
||||
snowflake.proxyType = proxyType
|
||||
snowflake.natType = natType
|
||||
snowflake.offerChannel = make(chan *ClientOffer)
|
||||
|
@ -205,7 +207,7 @@ func proxyPolls(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
|
|||
return
|
||||
}
|
||||
|
||||
sid, proxyType, natType, err := messages.DecodePollRequest(body)
|
||||
sid, proxyType, natType, clients, err := messages.DecodePollRequest(body)
|
||||
if err != nil {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
return
|
||||
|
@ -222,7 +224,7 @@ func proxyPolls(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
|
|||
}
|
||||
|
||||
// Wait for a client to avail an offer to the snowflake, or timeout if nil.
|
||||
offer := ctx.RequestOffer(sid, proxyType, natType)
|
||||
offer := ctx.RequestOffer(sid, proxyType, natType, clients)
|
||||
var b []byte
|
||||
if nil == offer {
|
||||
ctx.metrics.lock.Lock()
|
||||
|
|
|
@ -32,7 +32,7 @@ func TestBroker(t *testing.T) {
|
|||
Convey("Adds Snowflake", func() {
|
||||
So(ctx.snowflakes.Len(), ShouldEqual, 0)
|
||||
So(len(ctx.idToSnowflake), ShouldEqual, 0)
|
||||
ctx.AddSnowflake("foo", "", NATUnrestricted)
|
||||
ctx.AddSnowflake("foo", "", NATUnrestricted, 0)
|
||||
So(ctx.snowflakes.Len(), ShouldEqual, 1)
|
||||
So(len(ctx.idToSnowflake), ShouldEqual, 1)
|
||||
})
|
||||
|
@ -59,7 +59,7 @@ func TestBroker(t *testing.T) {
|
|||
Convey("Request an offer from the Snowflake Heap", func() {
|
||||
done := make(chan *ClientOffer)
|
||||
go func() {
|
||||
offer := ctx.RequestOffer("test", "", NATUnrestricted)
|
||||
offer := ctx.RequestOffer("test", "", NATUnrestricted, 0)
|
||||
done <- offer
|
||||
}()
|
||||
request := <-ctx.proxyPolls
|
||||
|
@ -84,7 +84,7 @@ func TestBroker(t *testing.T) {
|
|||
Convey("with a proxy answer if available.", func() {
|
||||
done := make(chan bool)
|
||||
// Prepare a fake proxy to respond with.
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted)
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted, 0)
|
||||
go func() {
|
||||
clientOffers(ctx, w, r)
|
||||
done <- true
|
||||
|
@ -102,7 +102,7 @@ func TestBroker(t *testing.T) {
|
|||
return
|
||||
}
|
||||
done := make(chan bool)
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted)
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted, 0)
|
||||
go func() {
|
||||
clientOffers(ctx, w, r)
|
||||
// Takes a few seconds here...
|
||||
|
@ -132,7 +132,7 @@ func TestBroker(t *testing.T) {
|
|||
Convey("with a proxy answer if available.", func() {
|
||||
done := make(chan bool)
|
||||
// Prepare a fake proxy to respond with.
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted)
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted, 0)
|
||||
go func() {
|
||||
clientOffers(ctx, w, r)
|
||||
done <- true
|
||||
|
@ -150,7 +150,7 @@ func TestBroker(t *testing.T) {
|
|||
return
|
||||
}
|
||||
done := make(chan bool)
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted)
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted, 0)
|
||||
go func() {
|
||||
clientOffers(ctx, w, r)
|
||||
// Takes a few seconds here...
|
||||
|
@ -201,7 +201,7 @@ func TestBroker(t *testing.T) {
|
|||
})
|
||||
|
||||
Convey("Responds to proxy answers...", func() {
|
||||
s := ctx.AddSnowflake("test", "", NATUnrestricted)
|
||||
s := ctx.AddSnowflake("test", "", NATUnrestricted, 0)
|
||||
w := httptest.NewRecorder()
|
||||
data := bytes.NewReader([]byte(`{"Version":"1.0","Sid":"test","Answer":"test"}`))
|
||||
|
||||
|
@ -314,7 +314,7 @@ func TestBroker(t *testing.T) {
|
|||
// Manually do the Broker goroutine action here for full control.
|
||||
p := <-ctx.proxyPolls
|
||||
So(p.id, ShouldEqual, "ymbcCMto7KHNGYlp")
|
||||
s := ctx.AddSnowflake(p.id, "", NATUnrestricted)
|
||||
s := ctx.AddSnowflake(p.id, "", NATUnrestricted, 0)
|
||||
go func() {
|
||||
offer := <-s.offerChannel
|
||||
p.offerChannel <- offer
|
||||
|
@ -593,7 +593,7 @@ func TestMetrics(t *testing.T) {
|
|||
So(err, ShouldBeNil)
|
||||
|
||||
// Prepare a fake proxy to respond with.
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted)
|
||||
snowflake := ctx.AddSnowflake("fake", "", NATUnrestricted, 0)
|
||||
go func() {
|
||||
clientOffers(ctx, w, r)
|
||||
done <- true
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue