├── .circleci └── config.yml ├── .dockerignore ├── Dockerfile ├── LICENSE ├── Makefile ├── README.md ├── circleci.sh ├── concurrentlimit.go ├── concurrentlimit_test.go ├── go.mod ├── go.sum ├── grpclimit ├── grpclimit.go └── grpclimit_test.go ├── limitserver └── limitserver.go ├── loadclient └── main.go ├── sleepymemory ├── sleepymemory.pb.go ├── sleepymemory.proto ├── sleepymemory_doc.go └── sleepymemory_grpc.pb.go └── sleepyserver └── sleepyserver.go /.circleci/config.yml: -------------------------------------------------------------------------------- 1 | version: 2.1 2 | 3 | workflows: 4 | version: 2 5 | build_and_test: 6 | jobs: 7 | - test 8 | - test_docker_image 9 | 10 | jobs: 11 | test: 12 | docker: 13 | - image: golang:1.20.1-bullseye 14 | steps: 15 | - checkout 16 | - run: 17 | name: run tests 18 | command: ./circleci.sh 19 | 20 | test_docker_image: 21 | docker: 22 | - image: cimg/base:edge 23 | steps: 24 | - checkout 25 | # Allow access to docker commands: https://circleci.com/docs/2.0/building-docker-images/ 26 | - setup_remote_docker 27 | - run: 28 | name: verify that the Dockerfile works 29 | command: docker build . --tag=dockerimage 30 | - run: 31 | name: verify that the built image mostly works 32 | command: docker run --rm -ti dockerimage --help 33 | -------------------------------------------------------------------------------- /.dockerignore: -------------------------------------------------------------------------------- 1 | # https://docs.docker.com/engine/reference/builder/#dockerignore-file 2 | **/*_test.go 3 | .circleci 4 | .dockerignore 5 | .git 6 | build 7 | circleci.sh 8 | Dockerfile 9 | LICENSE 10 | README.md 11 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # Go build image: separate downloading dependencies from build for incremental builds 2 | FROM golang:1.20.1-bullseye AS go_dep_downloader 3 | WORKDIR concurrentlimit 4 | COPY go.mod . 5 | COPY go.sum . 6 | RUN go mod download -x 7 | 8 | # Go build image: separate downloading dependencies from build for incremental builds 9 | FROM go_dep_downloader AS go_builder 10 | COPY . . 11 | RUN CGO_ENABLED=0 go install -v ./sleepyserver 12 | 13 | FROM gcr.io/distroless/static-debian11:nonroot AS sleepyserver 14 | COPY --from=go_builder /go/bin/sleepyserver / 15 | ENTRYPOINT ["/sleepyserver"] 16 | CMD ["--httpAddr=:8080", "--grpcAddr=:8081"] 17 | EXPOSE 8080 18 | EXPOSE 8081 19 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2020, Evan Jones. All rights reserved. 2 | 3 | Redistribution and use in source and binary forms, with or without 4 | modification, are permitted provided that the following conditions are 5 | met: 6 | 7 | * Redistributions of source code must retain the above copyright 8 | notice, this list of conditions and the following disclaimer. 9 | * Redistributions in binary form must reproduce the above 10 | copyright notice, this list of conditions and the following disclaimer 11 | in the documentation and/or other materials provided with the 12 | distribution. 13 | * Neither the name of the copyright holder nor the names of its 14 | contributors may be used to endorse or promote products derived from 15 | this software without specific prior written permission. 16 | 17 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 18 | "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 19 | LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 20 | A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 21 | OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 22 | SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 23 | LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 24 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 25 | THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 26 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | # Makefile for updating protocol buffer definitions and downloading required tools 2 | BUILD_DIR:=build 3 | PROTOC:=$(BUILD_DIR)/bin/protoc 4 | PROTOC_GEN_GO:=$(BUILD_DIR)/protoc-gen-go 5 | PROTOC_GEN_GO_GRPC:=$(BUILD_DIR)/protoc-gen-go-grpc 6 | 7 | sleepymemory/sleepymemory.pb.go: sleepymemory/sleepymemory.proto $(PROTOC) $(PROTOC_GEN_GO) $(PROTOC_GEN_GO_GRPC) 8 | $(PROTOC) --plugin=$(PROTOC_GEN_GO) --plugin=$(PROTOC_GEN_GO_GRPC) \ 9 | --go_out=paths=source_relative:. \ 10 | --go-grpc_out=paths=source_relative:. \ 11 | $< 12 | 13 | # download protoc to a temporary tools directory 14 | $(PROTOC): $(BUILD_DIR)/getprotoc | $(BUILD_DIR) 15 | $(BUILD_DIR)/getprotoc --outputDir=$(BUILD_DIR) 16 | 17 | $(BUILD_DIR)/getprotoc: | $(BUILD_DIR) 18 | GOBIN=$(realpath $(BUILD_DIR)) go install github.com/evanj/hacks/getprotoc@latest 19 | 20 | # go install uses the version of protoc-gen-go specified by go.mod ... I think 21 | $(PROTOC_GEN_GO): go.mod | $(BUILD_DIR) 22 | GOBIN=$(realpath $(BUILD_DIR)) go install google.golang.org/protobuf/cmd/protoc-gen-go 23 | 24 | # manually specified version since we don't import this from code anywhere 25 | # TODO: Import this from some tool so it gets updated with go get? 26 | $(PROTOC_GEN_GO_GRPC): go.mod | $(BUILD_DIR) 27 | GOBIN=$(realpath $(BUILD_DIR)) go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@v1.2.0 28 | 29 | $(BUILD_DIR): 30 | mkdir -p $@ 31 | 32 | clean: 33 | $(RM) -r $(BUILD_DIR) 34 | 35 | docker: 36 | docker build . --tag=gcr.io/networkping/sleepyserver:$(shell date '+%Y%m%d')-$(shell git rev-parse --short=10 HEAD) 37 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Concurrent request/connection limits for Go servers 2 | 3 | Each connection and request that a server is processing takes memory. If you have too many, your server will run out of memory and crash. To make servers robust, you must limit the amount of concurrent work that it accepts, so it at least serves some requests during overload scenarios, rather than serving none. This package provides some APIs to make this easier, and some tools I used to test this. 4 | 5 | For a really robust server, you should do the following, in rough priority order: 6 | 7 | * Limit concurrent executing requests to limit memory. 8 | * Limit concurrent connections to limit memory. In particular, Go gRPC connections are very expensive (~230 kiB versus about ~40 kiB per HTTP server connection, particularly when connections are opening/closing rapidly). 9 | * Close connections/requests that are too slow or idle, since they are wasting resources. 10 | * Make clients well-behaved so they reduce their request rate on error, or stop entirely (exponential backoff, back pressure, circuit breakers). The gRPC `MaxConcurrentStreams` setting can help here. 11 | 12 | 13 | # Possible future improvements to this code 14 | 15 | * *gRPC streaming requests*: The `grpclimit` package currently only limits unary requests. 16 | 17 | * *Faster implementation*: This uses a single sync.Mutex. It works well for ~10000 requests/second on 8 CPUs, but can be a bottleneck for extremely low-latency requests or high-CPU servers. Some sort of sharded counter, or something crazy like https://github.com/jonhoo/drwmutex would be more efficient. 18 | 19 | * *Blocking/queuing*: This package currently rejects requests when over the limit. It probably would be better to queue requests for some period of time. This can cause fewer retries when there are short overload bursts. It also means that poorly behaved clients that retry too quickly will retry less often, which may ultimately be better. There are also choices here about LIFO versus FIFO, drop head versus drop tail. See my previous investigation: https://www.evanjones.ca/prevent-server-overload.html 20 | 21 | * *Multiple buckets of limits*: Health checks, statistics, or other cheap requests should have much higher limits than expensive requests. It is possible this should be configurable. 22 | 23 | * *Aggressively close idle connections on overload*: This package sets idle timeouts on connections to attempt to avoid lots of idle clients starving busy clients. It would be nice if this policy triggered on overload. If we are at the connection limit, we should aggressively close idle connections. If we are not, then we should not care. 24 | 25 | 26 | ## Running the server with limited memory and Docker 27 | 28 | ``` 29 | docker build . --tag=sleepyserver 30 | docker run -p 127.0.0.1:8080:8080 -p 127.0.0.1:8081:8081 --rm -ti --memory=128m --memory-swap=128m sleepyserver 31 | ``` 32 | 33 | ## To monitor in another terminal: 34 | 35 | * `docker stats` 36 | * `curl http://localhost:8080/stats` 37 | 38 | ## High memory per request 39 | 40 | This client makes requests that use 1 MiB/request. Using 80 concurrent clients reliably blows up the server very quickly. Adding the concurrent rate limiter --concurrentRequests=40 fixes it. 41 | 42 | ``` 43 | ulimit -n 10000 44 | # HTTP 45 | go run ./loadclient/main.go --httpTarget=http://localhost:8080/ --concurrent=80 --sleep=3s --waste=1048576 --duration=2m 46 | # gRPC 47 | go run ./loadclient/main.go --grpcTarget=localhost:8081 --concurrent=80 --sleep=3s --waste=1048576 --duration=2m 48 | ``` 49 | 50 | 51 | ## Low memory per request (lots of idle requests) 52 | 53 | This client makes requests that basically do nothing except use idle connections. 54 | 55 | ``` 56 | ulimit -n 10000 57 | # HTTP 58 | go run ./loadclient/main.go --httpTarget=http://localhost:8080/ --concurrent=5000 --sleep=20s --duration=2m 59 | # gRPC 60 | go run ./loadclient/main.go --grpcTarget=localhost:8081 --concurrent=5000 --sleep=20s --duration=2m 61 | ``` 62 | 63 | With HTTP and a docker memory limit of 128 MiB, on my machine 3000 concurrent connections seems to "work" but is dangerously close to the limit. Running the test a few times in a row seems to kill it. It seems like closing and re-opening connections causes an increase in memory usage. The gRPC test fails at a lower connection count (around 1000), so those connections are MUCH more memory expensive than HTTP connections. 64 | 65 | * 3000-3100 works but unreliably 66 | * 3200 works for a while but dies 67 | * 3500 connections dies after a few minutes 68 | * 3800 connections reliably dies 69 | 70 | Using a concurrent request limit does NOT solve the problem, even with --concurrentRequests=40: There are simply too many connections and too much goroutine/connection overhead. To fix this, we need to reject new connections using --concurrentConnections=80. 71 | 72 | 73 | ## gRPC MaxConcurrentStreams 74 | 75 | This limits the number of concurrent streams *per-client connection*, so this doesn't fix overload by itself. For example, setting it to 40, and using the "high memory" client above still blows through the limit. With the `--shareGRPC` client, this will protect it. With this option, the server communicates the limit back to the client, which means the client will block and slow down its rate of requests (back-pressure). It is still useful, but does not protect the server's resources appropriately from "worst case" scenarios. 76 | -------------------------------------------------------------------------------- /circleci.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Runs checks on CircleCI 3 | 4 | set -euf -o pipefail 5 | 6 | # echo commands 7 | set -x 8 | 9 | # Ensure protocol buffer definitions are up to date 10 | make 11 | 12 | # Run tests 13 | go test -count=2 -shuffle=on -race ./... 14 | 15 | # go test only checks some vet warnings; check all 16 | go vet ./... 17 | 18 | go install honnef.co/go/tools/cmd/staticcheck@latest 19 | staticcheck --checks=all ./... 20 | 21 | go fmt ./... 22 | # require that we use go mod tidy. TODO: there must be an easier way? 23 | go mod tidy 24 | CHANGED=$(git status --porcelain --untracked-files=no) 25 | if [ -n "${CHANGED}" ]; then 26 | echo "ERROR files were changed:" > /dev/stderr 27 | echo "$CHANGED" > /dev/stderr 28 | exit 10 29 | fi 30 | -------------------------------------------------------------------------------- /concurrentlimit.go: -------------------------------------------------------------------------------- 1 | // Package concurrentlimit limits the number of concurrent requests to a Go HTTP or gRPC server. 2 | package concurrentlimit 3 | 4 | import ( 5 | "errors" 6 | "fmt" 7 | "log" 8 | "net" 9 | "net/http" 10 | "sync" 11 | "time" 12 | 13 | "golang.org/x/net/netutil" 14 | ) 15 | 16 | // ErrLimited is returned by Limiter when the concurrent operation limit is exceeded. 17 | var ErrLimited = errors.New("exceeded max concurrent operations limit") 18 | 19 | // This should be set longer than what upstream clients/load balancers will use to avoid 20 | // a "connection race" where the client sends a request at the same time the server is closing 21 | // it. This can cause errors that may not be retriable. This is the value recommended by Google 22 | // Cloud: https://cloud.google.com/load-balancing/docs/https#timeouts_and_retries 23 | const httpIdleTimeout = 620 * time.Second 24 | const httpReadHeaderTimeout = time.Minute 25 | 26 | // Limiter limits the number of concurrent operations that can be processed. 27 | type Limiter interface { 28 | // Start begins a new operation. It returns a completion function that must be called when the 29 | // operation completes, or it returns ErrLimited if no more concurrent operations are allowed. 30 | // This should be called as: 31 | // 32 | // end, err := limiter.Start() 33 | // if err != nil { 34 | // // Handle ErrLimited 35 | // defer end() 36 | Start() (func(), error) 37 | } 38 | 39 | // NoLimit returns a Limiter that permits an unlimited number of operations. 40 | func NoLimit() Limiter { 41 | return &nilLimiter{} 42 | } 43 | 44 | type nilLimiter struct{} 45 | 46 | func doNothing() {} 47 | 48 | func (n *nilLimiter) Start() (func(), error) { 49 | return doNothing, nil 50 | } 51 | 52 | // New returns a Limiter that will only permit limit concurrent operations. It will panic if 53 | // limit is < 0. 54 | func New(limit int) Limiter { 55 | if limit <= 0 { 56 | panic(fmt.Sprintf("limit must be > 0: %d", limit)) 57 | } 58 | return &syncLimiter{sync.Mutex{}, limit, 0} 59 | } 60 | 61 | type syncLimiter struct { 62 | mu sync.Mutex 63 | max int 64 | current int 65 | } 66 | 67 | func (s *syncLimiter) Start() (func(), error) { 68 | s.mu.Lock() 69 | defer s.mu.Unlock() 70 | 71 | next := s.current + 1 72 | if next > s.max { 73 | return nil, ErrLimited 74 | } 75 | s.current = next 76 | 77 | // TODO: Return a closure that can only be called once? More expensive but harder to abuse. 78 | // Maybe think about a "debug mode" that enables this sort of check? 79 | return s.end, nil 80 | } 81 | 82 | func (s *syncLimiter) end() { 83 | s.mu.Lock() 84 | s.current-- 85 | if s.current < 0 { 86 | panic("bug: mismatched calls to start/end") 87 | } 88 | s.mu.Unlock() 89 | } 90 | 91 | // ListenAndServe listens for HTTP requests with a limited number of concurrent requests 92 | // and connections. This helps avoid running out of memory during overload situations. 93 | // Both requestLimit and connectionLimit must be > 0, and connectionLimit must be 94 | // >= requestLimit. A reasonable defalt is to set the connectionLimit to double the request limit, 95 | // which assumes that processing each request requires more memory than a raw connection, and that 96 | // keeping some idle connections is useful. This modifies srv.Handler with another handler that 97 | // implements the limit. 98 | // 99 | // This also sets the server's ReadHeaderTimeout and IdleTimeout to a reasonable default if they 100 | // are not set, which is an attempt to avoid idle or slow connections using all connections. 101 | func ListenAndServe(srv *http.Server, requestLimit int, connectionLimit int) error { 102 | limitedListener, err := limitListenerForServer(srv, requestLimit, connectionLimit) 103 | if err != nil { 104 | return err 105 | } 106 | 107 | return srv.Serve(limitedListener) 108 | } 109 | 110 | func limitListenerForServer(srv *http.Server, requestLimit int, connectionLimit int) (net.Listener, error) { 111 | if requestLimit <= 0 { 112 | return nil, fmt.Errorf("ListenAndServe: requestLimit=%d must be > 0", requestLimit) 113 | } 114 | if connectionLimit < requestLimit { 115 | return nil, fmt.Errorf("ListenAndServe: connectionLimit=%d must be >= requestLimit=%d", 116 | connectionLimit, requestLimit) 117 | } 118 | 119 | // prevent idle/slow connections using all available connections. See also: 120 | // https://blog.gopheracademy.com/advent-2016/exposing-go-on-the-internet/ 121 | if srv.ReadHeaderTimeout <= 0 { 122 | srv.ReadHeaderTimeout = httpReadHeaderTimeout 123 | } 124 | if srv.IdleTimeout <= 0 { 125 | srv.IdleTimeout = httpIdleTimeout 126 | } 127 | 128 | // configure the request limit 129 | limiter := New(requestLimit) 130 | srv.Handler = Handler(limiter, srv.Handler) 131 | 132 | return Listen("tcp", srv.Addr, connectionLimit) 133 | } 134 | 135 | // ListenAndServeTLS listens for HTTP requests with a limited number of concurrent requests 136 | // and connections. See the documentation for ListenAndServe for details. 137 | func ListenAndServeTLS( 138 | srv *http.Server, certFile string, keyFile string, requestLimit int, connectionLimit int, 139 | ) error { 140 | limitedListener, err := limitListenerForServer(srv, requestLimit, connectionLimit) 141 | if err != nil { 142 | return err 143 | } 144 | 145 | return srv.ServeTLS(limitedListener, certFile, keyFile) 146 | } 147 | 148 | // Listen wraps net.Listen with netutil.LimitListener to limit concurrent connections. 149 | func Listen(network string, address string, connectionLimit int) (net.Listener, error) { 150 | unlimitedListener, err := net.Listen(network, address) 151 | if err != nil { 152 | return nil, err 153 | } 154 | return netutil.LimitListener(unlimitedListener, connectionLimit), nil 155 | } 156 | 157 | // Handler returns an http.Handler that uses limiter to only permit a limited number of concurrent 158 | // requests to be processed. 159 | func Handler(limiter Limiter, handler http.Handler) http.Handler { 160 | return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 161 | end, err := limiter.Start() 162 | if err == ErrLimited { 163 | http.Error(w, err.Error(), http.StatusTooManyRequests) 164 | return 165 | } 166 | if err != nil { 167 | // this should not happen, but if it does return a very generic 500 error 168 | log.Println("concurrentlimit.Handler BUG: unexpected error: " + err.Error()) 169 | http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) 170 | return 171 | } 172 | 173 | // permitted: start the operation and end it 174 | handler.ServeHTTP(w, r) 175 | end() 176 | }) 177 | } 178 | -------------------------------------------------------------------------------- /concurrentlimit_test.go: -------------------------------------------------------------------------------- 1 | package concurrentlimit 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "net" 7 | "net/http" 8 | "strconv" 9 | "sync" 10 | "syscall" 11 | "testing" 12 | "time" 13 | ) 14 | 15 | func TestNoLimit(t *testing.T) { 16 | limiter := NoLimit() 17 | 18 | endFuncs := []func(){} 19 | for i := 0; i < 10000; i++ { 20 | end, err := limiter.Start() 21 | if err != nil { 22 | t.Fatal("NoLimit should never return an error") 23 | } 24 | endFuncs = append(endFuncs, end) 25 | } 26 | 27 | // calling all the end functions should work 28 | for _, end := range endFuncs { 29 | end() 30 | } 31 | } 32 | 33 | func TestLimiterRace(t *testing.T) { 34 | const permitted = 100 35 | limiter := New(permitted) 36 | 37 | // start the limiter in separate goroutines so hopefully the race detector can find bugs 38 | var wg sync.WaitGroup 39 | endFuncs := make(chan func(), permitted) 40 | for i := 0; i < permitted; i++ { 41 | wg.Add(1) 42 | go func() { 43 | defer wg.Done() 44 | 45 | end, err := limiter.Start() 46 | if err != nil { 47 | t.Error("Limiter must allow the first N calls", err) 48 | } 49 | endFuncs <- end 50 | }() 51 | } 52 | wg.Wait() 53 | 54 | // the next calls must fail 55 | for i := 0; i < 5; i++ { 56 | end, err := limiter.Start() 57 | if !(end == nil && err == ErrLimited) { 58 | t.Fatalf("Limiter must block calls after the first N calls: %p %#v", end, err) 59 | } 60 | } 61 | 62 | // Call one end function: the next Start call should work 63 | end := <-endFuncs 64 | end() 65 | end, err := limiter.Start() 66 | if !(end != nil && err == nil) { 67 | t.Fatal("The next call must succeed after end is called") 68 | } 69 | endFuncs <- end 70 | close(endFuncs) 71 | 72 | // calling all the end functions should work 73 | for end := range endFuncs { 74 | end() 75 | } 76 | } 77 | 78 | // Block HTTP requests until unblock is closed 79 | type blockForConcurrent struct { 80 | unblock chan struct{} 81 | } 82 | 83 | func (b *blockForConcurrent) ServeHTTP(w http.ResponseWriter, r *http.Request) { 84 | <-b.unblock 85 | } 86 | 87 | func TestHTTP(t *testing.T) { 88 | // set up a rate limited HTTP server 89 | const permitted = 3 90 | 91 | // pick a random port that should be available 92 | listener, err := net.Listen("tcp", "localhost:0") 93 | if err != nil { 94 | t.Fatal(err) 95 | } 96 | port := listener.Addr().(*net.TCPAddr).Port 97 | err = listener.Close() 98 | if err != nil { 99 | t.Fatal(err) 100 | } 101 | httpAddr := "localhost:" + strconv.Itoa(port) 102 | 103 | // start the server 104 | handler := &blockForConcurrent{make(chan struct{})} 105 | testServer := &http.Server{ 106 | Addr: httpAddr, 107 | Handler: handler, 108 | } 109 | go func() { 110 | // must allow more connections than requests, otherwise it waits for the connection to close 111 | err := ListenAndServe(testServer, permitted, permitted*2) 112 | if err != http.ErrServerClosed { 113 | t.Error("expected HTTP server to be shutdown; err:", err) 114 | } 115 | }() 116 | defer testServer.Shutdown(context.Background()) 117 | 118 | responses := make(chan int) 119 | for i := 0; i < permitted+1; i++ { 120 | go func() { 121 | const attempts = 3 122 | for i := 0; i < attempts; i++ { 123 | resp, err := http.Get("http://" + httpAddr) 124 | if err != nil { 125 | var syscallErr syscall.Errno 126 | if errors.As(err, &syscallErr) && syscallErr == syscall.ECONNREFUSED { 127 | // race with the server starting up: try again 128 | time.Sleep(10 * time.Millisecond) 129 | continue 130 | } 131 | close(responses) 132 | t.Error(err) 133 | 134 | } 135 | resp.Body.Close() 136 | responses <- resp.StatusCode 137 | return 138 | } 139 | t.Error("failed after too many attempts") 140 | }() 141 | } 142 | 143 | okCount := 0 144 | rateLimitedCount := 0 145 | for i := 0; i < permitted+1; i++ { 146 | response := <-responses 147 | if i == 0 { 148 | // unblock the handlers on the first response, no matter what it is 149 | close(handler.unblock) 150 | } 151 | 152 | if response == http.StatusOK { 153 | okCount++ 154 | } else if response == http.StatusTooManyRequests { 155 | rateLimitedCount++ 156 | } else { 157 | t.Fatal("unexpected HTTP status code:", response) 158 | } 159 | } 160 | 161 | if !(okCount == permitted && rateLimitedCount == 1) { 162 | t.Error("unexpected OK and rate limited response counts:", okCount, rateLimitedCount) 163 | } 164 | } 165 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/evanj/concurrentlimit 2 | 3 | go 1.20 4 | 5 | require ( 6 | golang.org/x/net v0.7.0 7 | google.golang.org/grpc v1.53.0 8 | google.golang.org/protobuf v1.28.1 9 | ) 10 | 11 | require ( 12 | github.com/golang/protobuf v1.5.2 // indirect 13 | golang.org/x/sys v0.5.0 // indirect 14 | golang.org/x/text v0.7.0 // indirect 15 | google.golang.org/genproto v0.0.0-20230216225411-c8e22ba71e44 // indirect 16 | ) 17 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= 2 | github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw= 3 | github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= 4 | github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= 5 | github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= 6 | golang.org/x/net v0.7.0 h1:rJrUqqhjsgNp7KqAIc25s9pZnjU7TUcSY7HcVZjdn1g= 7 | golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= 8 | golang.org/x/sys v0.5.0 h1:MUK/U/4lj1t1oPg0HfuXDN/Z1wv31ZJ/YcPiGccS4DU= 9 | golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 10 | golang.org/x/text v0.7.0 h1:4BRB4x83lYWy72KwLD/qYDuTu7q9PjSagHvijDw7cLo= 11 | golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= 12 | golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 13 | google.golang.org/genproto v0.0.0-20230216225411-c8e22ba71e44 h1:EfLuoKW5WfkgVdDy7dTK8qSbH37AX5mj/MFh+bGPz14= 14 | google.golang.org/genproto v0.0.0-20230216225411-c8e22ba71e44/go.mod h1:8B0gmkoRebU8ukX6HP+4wrVQUY1+6PkQ44BSyIlflHA= 15 | google.golang.org/grpc v1.53.0 h1:LAv2ds7cmFV/XTS3XG1NneeENYrXGmorPxsBbptIjNc= 16 | google.golang.org/grpc v1.53.0/go.mod h1:OnIrk0ipVdj4N5d9IUoFUx72/VlD7+jUsHwZgwSMQpw= 17 | google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= 18 | google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= 19 | google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w= 20 | google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= 21 | -------------------------------------------------------------------------------- /grpclimit/grpclimit.go: -------------------------------------------------------------------------------- 1 | // Package grpclimit limits the number of concurrent requests and concurrent connections to a gRPC 2 | // server to ensure that it does not run out of memory during overload scenarios. 3 | package grpclimit 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | "time" 9 | 10 | "github.com/evanj/concurrentlimit" 11 | "google.golang.org/grpc" 12 | "google.golang.org/grpc/codes" 13 | "google.golang.org/grpc/keepalive" 14 | "google.golang.org/grpc/status" 15 | ) 16 | 17 | // ResourceExhausted seems slightly better than Unavailable, since 18 | // is wrong. Other examples: 19 | // 20 | // https://github.com/grpc/grpc/blob/master/doc/statuscodes.md: "Server temporarily out of 21 | // resources (e.g., Flow-control resource limits reached)" 22 | // 23 | // https://cloud.google.com/apis/design/errors: "Either out of resource quota or reaching rate 24 | // limiting" 25 | const rateLimitStatus = codes.ResourceExhausted 26 | 27 | // Set to the value recommended by the Google Cloud Load Balancer: 28 | // https://cloud.google.com/load-balancing/docs/https#timeouts_and_retries 29 | const idleConnectionTimeout = 620 * time.Second 30 | const keepaliveTimeout = time.Minute 31 | 32 | // NewServer creates a grpc.Server and net.Listener that supports a limited number of concurrent 33 | // requests. It sets the MaxConcurrentStreams option to the same value, which will cause 34 | // requests to block on the client if a single client sends too many requests. You should also use 35 | // Serve() with this server to protect against too many idle connections. 36 | // 37 | // NOTE: options must not contain any interceptors, since this function relies on adding our 38 | // own interceptor to limit the requests. Use NewServerWithInterceptors if you need interceptors. 39 | // 40 | // TODO: Implement stream interceptors 41 | func NewServer( 42 | addr string, requestLimit int, options ...grpc.ServerOption, 43 | ) (*grpc.Server, error) { 44 | return NewServerWithInterceptors(addr, requestLimit, nil, options...) 45 | } 46 | 47 | // NewServerWithInterceptors is a version of NewServer that permits customizing the interceptors. 48 | // The passed in interceptor will be called after the operation limiter permits the request. See 49 | // NewServer's documentation for the remaining details. 50 | func NewServerWithInterceptors( 51 | addr string, requestLimit int, unaryInterceptor grpc.UnaryServerInterceptor, 52 | options ...grpc.ServerOption, 53 | ) (*grpc.Server, error) { 54 | if requestLimit <= 0 { 55 | return nil, fmt.Errorf("NewServer: requestLimit=%d must be > 0", requestLimit) 56 | } 57 | 58 | requestLimiter := concurrentlimit.New(requestLimit) 59 | limitedUnaryInterceptorChain := UnaryInterceptor(requestLimiter, unaryInterceptor) 60 | 61 | options = append(options, grpc.MaxConcurrentStreams(uint32(requestLimit))) 62 | options = append(options, grpc.UnaryInterceptor(limitedUnaryInterceptorChain)) 63 | options = append(options, grpc.KeepaliveParams(keepalive.ServerParameters{ 64 | MaxConnectionIdle: idleConnectionTimeout, 65 | Time: keepaliveTimeout, 66 | })) 67 | server := grpc.NewServer(options...) 68 | 69 | return server, nil 70 | } 71 | 72 | // Serve listens on addr but only accepts a maximum of connectionLimit conenctions at one 73 | // time to limit memory usage. New connections will block in the kernel. This returns when 74 | // grpc.Server.Serve would normally return. 75 | func Serve(server *grpc.Server, addr string, connectionLimit int) error { 76 | if connectionLimit <= 0 { 77 | return fmt.Errorf("NewServer: connectionLimit=%d must be >= 0", connectionLimit) 78 | } 79 | 80 | listener, err := concurrentlimit.Listen("tcp", addr, connectionLimit) 81 | if err != nil { 82 | return err 83 | } 84 | return server.Serve(listener) 85 | } 86 | 87 | // UnaryInterceptor returns a grpc.UnaryServerInterceptor that uses limiter to limit the 88 | // concurrent requests. It will return codes.ResourceExhausted if the limiter rejects an operation. 89 | // If next is not nil, it will be called to chain the request handlers. If it is nil, this will 90 | // invoke the operation directly. 91 | func UnaryInterceptor(limiter concurrentlimit.Limiter, next grpc.UnaryServerInterceptor) grpc.UnaryServerInterceptor { 92 | return func( 93 | ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler, 94 | ) (interface{}, error) { 95 | end, err := limiter.Start() 96 | if err == concurrentlimit.ErrLimited { 97 | return nil, status.Error(rateLimitStatus, err.Error()) 98 | } 99 | if err != nil { 100 | return nil, err 101 | } 102 | defer end() 103 | 104 | if next != nil { 105 | return next(ctx, req, info, handler) 106 | } 107 | return handler(ctx, req) 108 | } 109 | } 110 | -------------------------------------------------------------------------------- /grpclimit/grpclimit_test.go: -------------------------------------------------------------------------------- 1 | package grpclimit 2 | 3 | import ( 4 | "context" 5 | "log" 6 | "net" 7 | "strconv" 8 | "testing" 9 | "time" 10 | 11 | "github.com/evanj/concurrentlimit/sleepymemory" 12 | "google.golang.org/grpc" 13 | "google.golang.org/grpc/codes" 14 | "google.golang.org/grpc/credentials/insecure" 15 | "google.golang.org/grpc/status" 16 | ) 17 | 18 | type blockSleeper struct { 19 | sleepymemory.UnimplementedSleeperServer 20 | unblock chan struct{} 21 | } 22 | 23 | func (b *blockSleeper) Sleep( 24 | ctx context.Context, request *sleepymemory.SleepRequest, 25 | ) (*sleepymemory.SleepResponse, error) { 26 | <-b.unblock 27 | 28 | return &sleepymemory.SleepResponse{}, nil 29 | } 30 | 31 | func TestGRPC(t *testing.T) { 32 | const permitted = 3 33 | 34 | // pick a random port that should be avaliable 35 | listener, err := net.Listen("tcp", "localhost:0") 36 | if err != nil { 37 | t.Fatal(err) 38 | } 39 | port := listener.Addr().(*net.TCPAddr).Port 40 | err = listener.Close() 41 | if err != nil { 42 | t.Fatal(err) 43 | } 44 | grpcAddr := "localhost:" + strconv.Itoa(port) 45 | 46 | // start the server 47 | grpcServer, err := NewServer(grpcAddr, permitted) 48 | if err != nil { 49 | panic(err) 50 | } 51 | handler := &blockSleeper{unblock: make(chan struct{})} 52 | sleepymemory.RegisterSleeperServer(grpcServer, handler) 53 | go func() { 54 | err = Serve(grpcServer, grpcAddr, permitted*2) 55 | if err != nil { 56 | t.Error(err) 57 | } 58 | }() 59 | defer grpcServer.GracefulStop() 60 | 61 | // make permitted + 1 requests on child goroutines 62 | responses := make(chan codes.Code) 63 | for i := 0; i < permitted+1; i++ { 64 | go func() { 65 | // need separate clients per goroutine otherwise the client back pressure prevents rejection 66 | // Sleep for at least 2 seconds due to gRPC's default connection backoff: 67 | // https://github.com/grpc/grpc/blob/master/doc/connection-backoff.md 68 | dialCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) 69 | conn, err := grpc.DialContext(dialCtx, grpcAddr, 70 | grpc.WithTransportCredentials(insecure.NewCredentials()), 71 | grpc.WithBlock()) 72 | cancel() 73 | if err != nil { 74 | log.Println("Dial:", err) 75 | close(responses) 76 | t.Error(err) 77 | } 78 | defer conn.Close() 79 | client := sleepymemory.NewSleeperClient(conn) 80 | 81 | _, err = client.Sleep(context.Background(), &sleepymemory.SleepRequest{}) 82 | responses <- status.Code(err) 83 | }() 84 | } 85 | 86 | okCount := 0 87 | rateLimitedCount := 0 88 | for i := 0; i < permitted+1; i++ { 89 | response := <-responses 90 | if i == 0 { 91 | // unblock the handlers on the first response, no matter what it is 92 | close(handler.unblock) 93 | } 94 | if response == codes.OK { 95 | okCount++ 96 | } else if response == codes.ResourceExhausted { 97 | rateLimitedCount++ 98 | } else { 99 | t.Fatal("status code:", response) 100 | } 101 | } 102 | 103 | if !(okCount == permitted && rateLimitedCount == 1) { 104 | t.Error("unexpected OK and rate limited response counts:", okCount, rateLimitedCount) 105 | } 106 | } 107 | -------------------------------------------------------------------------------- /limitserver/limitserver.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "fmt" 7 | "log" 8 | "net/http" 9 | "runtime" 10 | "strconv" 11 | "sync" 12 | "time" 13 | 14 | "github.com/evanj/concurrentlimit" 15 | "github.com/evanj/concurrentlimit/grpclimit" 16 | "github.com/evanj/concurrentlimit/sleepymemory" 17 | "google.golang.org/grpc/codes" 18 | "google.golang.org/grpc/status" 19 | "google.golang.org/protobuf/types/known/durationpb" 20 | ) 21 | 22 | const sleepHTTPKey = "sleep" 23 | const wasteHTTPKey = "waste" 24 | 25 | type server struct { 26 | sleepymemory.UnimplementedSleeperServer 27 | logger concurrentMaxLogger 28 | } 29 | 30 | func (s *server) rawRootHandler(w http.ResponseWriter, r *http.Request) { 31 | if r.URL.Path != "/" { 32 | http.NotFound(w, r) 33 | return 34 | } 35 | if r.Method != http.MethodGet { 36 | http.Error(w, "only GET is supported", http.StatusMethodNotAllowed) 37 | return 38 | } 39 | 40 | err := s.rootHandler(w, r) 41 | if err != nil { 42 | statusCode := http.StatusInternalServerError 43 | if err == concurrentlimit.ErrLimited { 44 | statusCode = http.StatusTooManyRequests 45 | } 46 | http.Error(w, err.Error(), statusCode) 47 | } 48 | } 49 | 50 | func humanBytes(bytes uint64) string { 51 | megabytes := float64(bytes) / float64(1024*1024) 52 | return fmt.Sprintf("%.1f", megabytes) 53 | } 54 | 55 | func (s *server) memstatsHandler(w http.ResponseWriter, r *http.Request) { 56 | stats := &runtime.MemStats{} 57 | runtime.ReadMemStats(stats) 58 | 59 | w.Header().Set("Content-Type", "text/plain;charset=utf-8") 60 | fmt.Fprintf(w, "total bytes of memory obtained from the OS Sys=%d %s\n", 61 | stats.Sys, humanBytes(stats.Sys)) 62 | fmt.Fprintf(w, "bytes of allocated heap objects HeapAlloc=%d %s\n", 63 | stats.HeapAlloc, humanBytes(stats.HeapAlloc)) 64 | } 65 | 66 | func (s *server) rootHandler(w http.ResponseWriter, r *http.Request) error { 67 | req := &sleepymemory.SleepRequest{} 68 | 69 | sleepValue := r.FormValue(sleepHTTPKey) 70 | if sleepValue != "" { 71 | var sleepDuration time.Duration 72 | // try to parse as integer seconds first 73 | seconds, err := strconv.Atoi(sleepValue) 74 | if err == nil { 75 | // SUCCESS! 76 | sleepDuration = time.Duration(seconds) * time.Second 77 | } else { 78 | // fall back to parsing duration, and return that error if it fails 79 | sleepDuration, err = time.ParseDuration(sleepValue) 80 | if err != nil { 81 | return err 82 | } 83 | } 84 | req.SleepDuration = durationpb.New(sleepDuration) 85 | } 86 | 87 | wasteValue := r.FormValue(wasteHTTPKey) 88 | if wasteValue != "" { 89 | bytes, err := strconv.Atoi(wasteValue) 90 | if err != nil { 91 | return err 92 | } 93 | req.WasteBytes = int64(bytes) 94 | } 95 | 96 | resp, err := s.sleepImplementation(r.Context(), req) 97 | if err != nil { 98 | return err 99 | } 100 | 101 | w.Header().Set("Content-Type", "text/plain;charset=utf-8") 102 | fmt.Fprintf(w, "slept for %s (pass ?sleep=x)\nwasted %d bytes (pass ?waste=y)\nignored response=%d\n", 103 | req.SleepDuration.String(), req.WasteBytes, resp.Ignored) 104 | return nil 105 | } 106 | 107 | func (s *server) Sleep(ctx context.Context, request *sleepymemory.SleepRequest) (*sleepymemory.SleepResponse, error) { 108 | resp, err := s.sleepImplementation(ctx, request) 109 | if err == concurrentlimit.ErrLimited { 110 | err = status.Error(codes.ResourceExhausted, err.Error()) 111 | } 112 | return resp, err 113 | } 114 | 115 | func (s *server) sleepImplementation(ctx context.Context, request *sleepymemory.SleepRequest) (*sleepymemory.SleepResponse, error) { 116 | // log max concurrent requests 117 | defer s.logger.start()() 118 | 119 | // waste memory and touch each page to ensure it is actually allocated 120 | wasteSlice := make([]byte, request.WasteBytes) 121 | const pageSize = 4096 122 | for i := 0; i < len(wasteSlice); i += pageSize { 123 | wasteSlice[i] = 0xff 124 | } 125 | 126 | var duration time.Duration 127 | if request.SleepDuration != nil { 128 | if err := request.SleepDuration.CheckValid(); err != nil { 129 | return nil, err 130 | } 131 | duration = request.SleepDuration.AsDuration() 132 | } 133 | // TODO: use ctx for cancellation 134 | time.Sleep(duration) 135 | 136 | // read some of the memory and return it so it doesn't get garbage collected 137 | total := 0 138 | for i := 0; i < len(wasteSlice); i += 10 * pageSize { 139 | total += int(wasteSlice[i]) 140 | } 141 | 142 | return &sleepymemory.SleepResponse{Ignored: int64(total)}, nil 143 | } 144 | 145 | type concurrentMaxLogger struct { 146 | mu sync.Mutex 147 | max int 148 | current int 149 | } 150 | 151 | // start records the start of a concurrent request that is terminated when func is called. 152 | func (c *concurrentMaxLogger) start() func() { 153 | c.mu.Lock() 154 | c.current++ 155 | if c.current > c.max { 156 | c.max = c.current 157 | log.Printf("new max requests=%d", c.max) 158 | } 159 | c.mu.Unlock() 160 | 161 | return c.end 162 | } 163 | 164 | func (c *concurrentMaxLogger) end() { 165 | c.mu.Lock() 166 | c.current-- 167 | if c.current < 0 { 168 | panic("bug: mismatched calls to startRequest/endRequest") 169 | } 170 | c.mu.Unlock() 171 | } 172 | 173 | func main() { 174 | httpAddr := flag.String("httpAddr", "localhost:8080", "Address to listen for HTTP requests") 175 | grpcAddr := flag.String("grpcAddr", "localhost:8081", "Address to listen for gRPC requests") 176 | concurrentRequests := flag.Int("concurrentRequests", 0, "Limits the number of concurrent requests") 177 | concurrentConnections := flag.Int("concurrentConnections", 0, "Limits the number of concurrent connections") 178 | flag.Parse() 179 | 180 | s := &server{logger: concurrentMaxLogger{}} 181 | 182 | mux := &http.ServeMux{} 183 | mux.HandleFunc("/", s.rawRootHandler) 184 | mux.HandleFunc("/stats", s.memstatsHandler) 185 | log.Printf("listening for HTTP on http://%s concurrentRequests=%d concurrentConnections=%d ...", 186 | *httpAddr, *concurrentRequests, *concurrentConnections) 187 | httpServer := &http.Server{ 188 | Addr: *httpAddr, 189 | Handler: mux, 190 | } 191 | 192 | go func() { 193 | err := concurrentlimit.ListenAndServe(httpServer, *concurrentRequests, *concurrentConnections) 194 | if err != nil { 195 | panic(err) 196 | } 197 | }() 198 | 199 | log.Printf("listening for gRPC on grpcAddr=%s concurrentRequests=%d concurrentConnections=%d ...", 200 | *grpcAddr, *concurrentRequests, *concurrentConnections) 201 | grpcServer, err := grpclimit.NewServer(*grpcAddr, *concurrentRequests) 202 | if err != nil { 203 | panic(err) 204 | } 205 | 206 | sleepymemory.RegisterSleeperServer(grpcServer, s) 207 | err = grpclimit.Serve(grpcServer, *grpcAddr, *concurrentConnections) 208 | if err != nil { 209 | panic(err) 210 | } 211 | } 212 | -------------------------------------------------------------------------------- /loadclient/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "flag" 7 | "fmt" 8 | "io" 9 | "log" 10 | "net/http" 11 | "strings" 12 | "time" 13 | 14 | "github.com/evanj/concurrentlimit/sleepymemory" 15 | "google.golang.org/grpc" 16 | "google.golang.org/grpc/codes" 17 | "google.golang.org/grpc/credentials/insecure" 18 | "google.golang.org/grpc/status" 19 | "google.golang.org/protobuf/types/known/durationpb" 20 | ) 21 | 22 | const grpcConnectTimeout = 30 * time.Second 23 | 24 | func sendRequestsGoroutine( 25 | done <-chan struct{}, totalRequestsChan chan<- int, sender requestSender, 26 | req *sleepymemory.SleepRequest, 27 | ) { 28 | // create a new sender for each goroutine 29 | sender = sender.clone() 30 | 31 | requestCount := 0 32 | sendLoop: 33 | for { 34 | // if done is closed, break out of the loop 35 | select { 36 | case <-done: 37 | break sendLoop 38 | default: 39 | } 40 | 41 | err := sender.send(req) 42 | if err != nil { 43 | if err == errRetry || err == context.DeadlineExceeded { 44 | // TODO: exponential backoff? 45 | time.Sleep(time.Second) 46 | continue 47 | } 48 | panic(err) 49 | } 50 | 51 | requestCount++ 52 | } 53 | totalRequestsChan <- requestCount 54 | } 55 | 56 | var errRetry = errors.New("retriable error") 57 | 58 | type requestSender interface { 59 | clone() requestSender 60 | send(req *sleepymemory.SleepRequest) error 61 | } 62 | 63 | type httpSender struct { 64 | client *http.Client 65 | baseURL string 66 | } 67 | 68 | func newHTTPSender(baseURL string) *httpSender { 69 | return &httpSender{ 70 | // create a separate client and transport so each goroutine uses a separate connection 71 | &http.Client{Transport: &http.Transport{}}, 72 | baseURL, 73 | } 74 | } 75 | 76 | func (h *httpSender) clone() requestSender { 77 | return newHTTPSender(h.baseURL) 78 | } 79 | 80 | func (h *httpSender) send(req *sleepymemory.SleepRequest) error { 81 | reqURL := fmt.Sprintf("%s?sleep=%d&waste=%d", 82 | h.baseURL, req.SleepDuration.Seconds, req.WasteBytes) 83 | 84 | resp, err := h.client.Get(reqURL) 85 | if err != nil { 86 | // docker's proxy is pretty unhappy with how we hit this with tons of connections concurrently 87 | // we also get "operation timed out" when we are limiting the number of connections 88 | if strings.Contains(err.Error(), "connection reset by peer") || 89 | strings.Contains(err.Error(), "write: broken pipe") || 90 | strings.Contains(err.Error(), "operation timed out") { 91 | return errRetry 92 | } 93 | return err 94 | } 95 | defer resp.Body.Close() 96 | 97 | // drain the body so the connection can be reused by keep alives 98 | _, err = io.Copy(io.Discard, resp.Body) 99 | if err != nil { 100 | return err 101 | } 102 | err = resp.Body.Close() 103 | if err != nil { 104 | return err 105 | } 106 | 107 | if resp.StatusCode == http.StatusTooManyRequests { 108 | // we were rate limited! Try again later 109 | return errRetry 110 | } else if resp.StatusCode != http.StatusOK { 111 | return errors.New("expected status ok: " + resp.Status) 112 | } 113 | return nil 114 | } 115 | 116 | type grpcSender struct { 117 | addr string 118 | client sleepymemory.SleeperClient 119 | } 120 | 121 | func newGRPCSender(addr string) *grpcSender { 122 | return &grpcSender{addr: addr} 123 | } 124 | 125 | func (g *grpcSender) clone() requestSender { 126 | // copies the client: if it was used it will share a connection 127 | cloned := *g 128 | return &cloned 129 | } 130 | 131 | func (g *grpcSender) send(req *sleepymemory.SleepRequest) error { 132 | if g.client == nil { 133 | dialCtx, cancel := context.WithTimeout(context.Background(), grpcConnectTimeout) 134 | conn, err := grpc.DialContext(dialCtx, g.addr, 135 | grpc.WithTransportCredentials(insecure.NewCredentials()), 136 | grpc.WithBlock()) 137 | cancel() 138 | if err != nil { 139 | return err 140 | } 141 | g.client = sleepymemory.NewSleeperClient(conn) 142 | } 143 | 144 | _, err := g.client.Sleep(context.Background(), req) 145 | if status.Code(err) == codes.ResourceExhausted { 146 | err = errRetry 147 | } 148 | return err 149 | } 150 | 151 | func main() { 152 | httpTarget := flag.String("httpTarget", "", "HTTP address to send requests to") 153 | grpcTarget := flag.String("grpcTarget", "", "HTTP address to send requests to") 154 | duration := flag.Duration("duration", time.Minute, "Duration to run the test") 155 | concurrent := flag.Int("concurrent", 1, "Number of concurrent client goroutines") 156 | sleep := flag.Duration("sleep", 0, "Time for the server to sleep handling a request") 157 | waste := flag.Int("waste", 0, "Bytes of memory the server should waste while handling a request") 158 | shareGRPC := flag.Bool("shareGRPC", false, "If set, the gRPC goroutines will share a single client") 159 | flag.Parse() 160 | 161 | req := &sleepymemory.SleepRequest{ 162 | SleepDuration: durationpb.New(*sleep), 163 | WasteBytes: int64(*waste), 164 | } 165 | 166 | var sender requestSender 167 | if *httpTarget != "" { 168 | log.Printf("sending HTTP requests to %s ...", *httpTarget) 169 | sender = newHTTPSender(*httpTarget) 170 | } else if *grpcTarget != "" { 171 | log.Printf("sending gRPC requests to %s ...", *grpcTarget) 172 | sender = newGRPCSender(*grpcTarget) 173 | if *shareGRPC { 174 | // make a request to create the client before we clone it so it will be shared 175 | log.Printf("sharing a single gRPC connection ...") 176 | err := sender.send(req) 177 | if err != nil { 178 | panic(err) 179 | } 180 | } 181 | } else { 182 | panic("specify --httpTarget or --grpcTarget") 183 | } 184 | 185 | log.Printf("sending requests for %s using %d client goroutines ...", 186 | duration.String(), *concurrent) 187 | done := make(chan struct{}) 188 | totalRequestsChan := make(chan int) 189 | for i := 0; i < *concurrent; i++ { 190 | go sendRequestsGoroutine(done, totalRequestsChan, sender, req) 191 | } 192 | 193 | time.Sleep(*duration) 194 | close(done) 195 | 196 | totalRequests := 0 197 | for i := 0; i < *concurrent; i++ { 198 | totalRequests += <-totalRequestsChan 199 | } 200 | close(totalRequestsChan) 201 | 202 | log.Printf("sent %d requests in %s using %d clients = %.3f reqs/sec", 203 | totalRequests, duration.String(), *concurrent, float64(totalRequests)/duration.Seconds()) 204 | } 205 | -------------------------------------------------------------------------------- /sleepymemory/sleepymemory.pb.go: -------------------------------------------------------------------------------- 1 | // Code generated by protoc-gen-go. DO NOT EDIT. 2 | // versions: 3 | // protoc-gen-go v1.28.1 4 | // protoc v4.22.0 5 | // source: sleepymemory/sleepymemory.proto 6 | 7 | package sleepymemory 8 | 9 | import ( 10 | protoreflect "google.golang.org/protobuf/reflect/protoreflect" 11 | protoimpl "google.golang.org/protobuf/runtime/protoimpl" 12 | durationpb "google.golang.org/protobuf/types/known/durationpb" 13 | reflect "reflect" 14 | sync "sync" 15 | ) 16 | 17 | const ( 18 | // Verify that this generated code is sufficiently up-to-date. 19 | _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) 20 | // Verify that runtime/protoimpl is sufficiently up-to-date. 21 | _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) 22 | ) 23 | 24 | type SleepRequest struct { 25 | state protoimpl.MessageState 26 | sizeCache protoimpl.SizeCache 27 | unknownFields protoimpl.UnknownFields 28 | 29 | // The duration the request will sleep. 30 | SleepDuration *durationpb.Duration `protobuf:"bytes,1,opt,name=sleep_duration,json=sleepDuration,proto3" json:"sleep_duration,omitempty"` 31 | // Bytes of memory that will be allocated before sleeping to simulate requests that use lots of memory. 32 | WasteBytes int64 `protobuf:"varint,2,opt,name=waste_bytes,json=wasteBytes,proto3" json:"waste_bytes,omitempty"` 33 | } 34 | 35 | func (x *SleepRequest) Reset() { 36 | *x = SleepRequest{} 37 | if protoimpl.UnsafeEnabled { 38 | mi := &file_sleepymemory_sleepymemory_proto_msgTypes[0] 39 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 40 | ms.StoreMessageInfo(mi) 41 | } 42 | } 43 | 44 | func (x *SleepRequest) String() string { 45 | return protoimpl.X.MessageStringOf(x) 46 | } 47 | 48 | func (*SleepRequest) ProtoMessage() {} 49 | 50 | func (x *SleepRequest) ProtoReflect() protoreflect.Message { 51 | mi := &file_sleepymemory_sleepymemory_proto_msgTypes[0] 52 | if protoimpl.UnsafeEnabled && x != nil { 53 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 54 | if ms.LoadMessageInfo() == nil { 55 | ms.StoreMessageInfo(mi) 56 | } 57 | return ms 58 | } 59 | return mi.MessageOf(x) 60 | } 61 | 62 | // Deprecated: Use SleepRequest.ProtoReflect.Descriptor instead. 63 | func (*SleepRequest) Descriptor() ([]byte, []int) { 64 | return file_sleepymemory_sleepymemory_proto_rawDescGZIP(), []int{0} 65 | } 66 | 67 | func (x *SleepRequest) GetSleepDuration() *durationpb.Duration { 68 | if x != nil { 69 | return x.SleepDuration 70 | } 71 | return nil 72 | } 73 | 74 | func (x *SleepRequest) GetWasteBytes() int64 { 75 | if x != nil { 76 | return x.WasteBytes 77 | } 78 | return 0 79 | } 80 | 81 | type SleepResponse struct { 82 | state protoimpl.MessageState 83 | sizeCache protoimpl.SizeCache 84 | unknownFields protoimpl.UnknownFields 85 | 86 | // The value is ignored but exists to prevent garbage collection freeing the waste slice early. 87 | Ignored int64 `protobuf:"varint,1,opt,name=ignored,proto3" json:"ignored,omitempty"` 88 | } 89 | 90 | func (x *SleepResponse) Reset() { 91 | *x = SleepResponse{} 92 | if protoimpl.UnsafeEnabled { 93 | mi := &file_sleepymemory_sleepymemory_proto_msgTypes[1] 94 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 95 | ms.StoreMessageInfo(mi) 96 | } 97 | } 98 | 99 | func (x *SleepResponse) String() string { 100 | return protoimpl.X.MessageStringOf(x) 101 | } 102 | 103 | func (*SleepResponse) ProtoMessage() {} 104 | 105 | func (x *SleepResponse) ProtoReflect() protoreflect.Message { 106 | mi := &file_sleepymemory_sleepymemory_proto_msgTypes[1] 107 | if protoimpl.UnsafeEnabled && x != nil { 108 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 109 | if ms.LoadMessageInfo() == nil { 110 | ms.StoreMessageInfo(mi) 111 | } 112 | return ms 113 | } 114 | return mi.MessageOf(x) 115 | } 116 | 117 | // Deprecated: Use SleepResponse.ProtoReflect.Descriptor instead. 118 | func (*SleepResponse) Descriptor() ([]byte, []int) { 119 | return file_sleepymemory_sleepymemory_proto_rawDescGZIP(), []int{1} 120 | } 121 | 122 | func (x *SleepResponse) GetIgnored() int64 { 123 | if x != nil { 124 | return x.Ignored 125 | } 126 | return 0 127 | } 128 | 129 | var File_sleepymemory_sleepymemory_proto protoreflect.FileDescriptor 130 | 131 | var file_sleepymemory_sleepymemory_proto_rawDesc = []byte{ 132 | 0x0a, 0x1f, 0x73, 0x6c, 0x65, 0x65, 0x70, 0x79, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x2f, 0x73, 133 | 0x6c, 0x65, 0x65, 0x70, 0x79, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x2e, 0x70, 0x72, 0x6f, 0x74, 134 | 0x6f, 0x12, 0x0c, 0x73, 0x6c, 0x65, 0x65, 0x70, 0x79, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x1a, 135 | 0x1e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 136 | 0x2f, 0x64, 0x75, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 137 | 0x71, 0x0a, 0x0c, 0x53, 0x6c, 0x65, 0x65, 0x70, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 138 | 0x40, 0x0a, 0x0e, 0x73, 0x6c, 0x65, 0x65, 0x70, 0x5f, 0x64, 0x75, 0x72, 0x61, 0x74, 0x69, 0x6f, 139 | 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x19, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 140 | 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x44, 0x75, 0x72, 0x61, 0x74, 0x69, 141 | 0x6f, 0x6e, 0x52, 0x0d, 0x73, 0x6c, 0x65, 0x65, 0x70, 0x44, 0x75, 0x72, 0x61, 0x74, 0x69, 0x6f, 142 | 0x6e, 0x12, 0x1f, 0x0a, 0x0b, 0x77, 0x61, 0x73, 0x74, 0x65, 0x5f, 0x62, 0x79, 0x74, 0x65, 0x73, 143 | 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x52, 0x0a, 0x77, 0x61, 0x73, 0x74, 0x65, 0x42, 0x79, 0x74, 144 | 0x65, 0x73, 0x22, 0x29, 0x0a, 0x0d, 0x53, 0x6c, 0x65, 0x65, 0x70, 0x52, 0x65, 0x73, 0x70, 0x6f, 145 | 0x6e, 0x73, 0x65, 0x12, 0x18, 0x0a, 0x07, 0x69, 0x67, 0x6e, 0x6f, 0x72, 0x65, 0x64, 0x18, 0x01, 146 | 0x20, 0x01, 0x28, 0x03, 0x52, 0x07, 0x69, 0x67, 0x6e, 0x6f, 0x72, 0x65, 0x64, 0x32, 0x4b, 0x0a, 147 | 0x07, 0x53, 0x6c, 0x65, 0x65, 0x70, 0x65, 0x72, 0x12, 0x40, 0x0a, 0x05, 0x53, 0x6c, 0x65, 0x65, 148 | 0x70, 0x12, 0x1a, 0x2e, 0x73, 0x6c, 0x65, 0x65, 0x70, 0x79, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 149 | 0x2e, 0x53, 0x6c, 0x65, 0x65, 0x70, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 150 | 0x73, 0x6c, 0x65, 0x65, 0x70, 0x79, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x2e, 0x53, 0x6c, 0x65, 151 | 0x65, 0x70, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x42, 0x2f, 0x5a, 0x2d, 0x67, 0x69, 152 | 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x65, 0x76, 0x61, 0x6e, 0x6a, 0x2f, 0x63, 153 | 0x6f, 0x6e, 0x63, 0x75, 0x72, 0x72, 0x65, 0x6e, 0x74, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x2f, 0x73, 154 | 0x6c, 0x65, 0x65, 0x70, 0x79, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x62, 0x06, 0x70, 0x72, 0x6f, 155 | 0x74, 0x6f, 0x33, 156 | } 157 | 158 | var ( 159 | file_sleepymemory_sleepymemory_proto_rawDescOnce sync.Once 160 | file_sleepymemory_sleepymemory_proto_rawDescData = file_sleepymemory_sleepymemory_proto_rawDesc 161 | ) 162 | 163 | func file_sleepymemory_sleepymemory_proto_rawDescGZIP() []byte { 164 | file_sleepymemory_sleepymemory_proto_rawDescOnce.Do(func() { 165 | file_sleepymemory_sleepymemory_proto_rawDescData = protoimpl.X.CompressGZIP(file_sleepymemory_sleepymemory_proto_rawDescData) 166 | }) 167 | return file_sleepymemory_sleepymemory_proto_rawDescData 168 | } 169 | 170 | var file_sleepymemory_sleepymemory_proto_msgTypes = make([]protoimpl.MessageInfo, 2) 171 | var file_sleepymemory_sleepymemory_proto_goTypes = []interface{}{ 172 | (*SleepRequest)(nil), // 0: sleepymemory.SleepRequest 173 | (*SleepResponse)(nil), // 1: sleepymemory.SleepResponse 174 | (*durationpb.Duration)(nil), // 2: google.protobuf.Duration 175 | } 176 | var file_sleepymemory_sleepymemory_proto_depIdxs = []int32{ 177 | 2, // 0: sleepymemory.SleepRequest.sleep_duration:type_name -> google.protobuf.Duration 178 | 0, // 1: sleepymemory.Sleeper.Sleep:input_type -> sleepymemory.SleepRequest 179 | 1, // 2: sleepymemory.Sleeper.Sleep:output_type -> sleepymemory.SleepResponse 180 | 2, // [2:3] is the sub-list for method output_type 181 | 1, // [1:2] is the sub-list for method input_type 182 | 1, // [1:1] is the sub-list for extension type_name 183 | 1, // [1:1] is the sub-list for extension extendee 184 | 0, // [0:1] is the sub-list for field type_name 185 | } 186 | 187 | func init() { file_sleepymemory_sleepymemory_proto_init() } 188 | func file_sleepymemory_sleepymemory_proto_init() { 189 | if File_sleepymemory_sleepymemory_proto != nil { 190 | return 191 | } 192 | if !protoimpl.UnsafeEnabled { 193 | file_sleepymemory_sleepymemory_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} { 194 | switch v := v.(*SleepRequest); i { 195 | case 0: 196 | return &v.state 197 | case 1: 198 | return &v.sizeCache 199 | case 2: 200 | return &v.unknownFields 201 | default: 202 | return nil 203 | } 204 | } 205 | file_sleepymemory_sleepymemory_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} { 206 | switch v := v.(*SleepResponse); i { 207 | case 0: 208 | return &v.state 209 | case 1: 210 | return &v.sizeCache 211 | case 2: 212 | return &v.unknownFields 213 | default: 214 | return nil 215 | } 216 | } 217 | } 218 | type x struct{} 219 | out := protoimpl.TypeBuilder{ 220 | File: protoimpl.DescBuilder{ 221 | GoPackagePath: reflect.TypeOf(x{}).PkgPath(), 222 | RawDescriptor: file_sleepymemory_sleepymemory_proto_rawDesc, 223 | NumEnums: 0, 224 | NumMessages: 2, 225 | NumExtensions: 0, 226 | NumServices: 1, 227 | }, 228 | GoTypes: file_sleepymemory_sleepymemory_proto_goTypes, 229 | DependencyIndexes: file_sleepymemory_sleepymemory_proto_depIdxs, 230 | MessageInfos: file_sleepymemory_sleepymemory_proto_msgTypes, 231 | }.Build() 232 | File_sleepymemory_sleepymemory_proto = out.File 233 | file_sleepymemory_sleepymemory_proto_rawDesc = nil 234 | file_sleepymemory_sleepymemory_proto_goTypes = nil 235 | file_sleepymemory_sleepymemory_proto_depIdxs = nil 236 | } 237 | -------------------------------------------------------------------------------- /sleepymemory/sleepymemory.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package sleepymemory; 4 | 5 | import "google/protobuf/duration.proto"; 6 | 7 | option go_package = "github.com/evanj/concurrentlimit/sleepymemory"; 8 | 9 | message SleepRequest { 10 | // The duration the request will sleep. 11 | google.protobuf.Duration sleep_duration = 1; 12 | 13 | // Bytes of memory that will be allocated before sleeping to simulate requests that use lots of memory. 14 | int64 waste_bytes = 2; 15 | } 16 | 17 | message SleepResponse { 18 | // The value is ignored but exists to prevent garbage collection freeing the waste slice early. 19 | int64 ignored = 1; 20 | } 21 | 22 | // Sleeper will sleep and waste memory to test concurrent requests and memory limits. 23 | service Sleeper { 24 | rpc Sleep (SleepRequest) returns (SleepResponse); 25 | } 26 | -------------------------------------------------------------------------------- /sleepymemory/sleepymemory_doc.go: -------------------------------------------------------------------------------- 1 | // Package sleepymemory contains protobuf messages. 2 | package sleepymemory 3 | -------------------------------------------------------------------------------- /sleepymemory/sleepymemory_grpc.pb.go: -------------------------------------------------------------------------------- 1 | // Code generated by protoc-gen-go-grpc. DO NOT EDIT. 2 | // versions: 3 | // - protoc-gen-go-grpc v1.2.0 4 | // - protoc v4.22.0 5 | // source: sleepymemory/sleepymemory.proto 6 | 7 | package sleepymemory 8 | 9 | import ( 10 | context "context" 11 | grpc "google.golang.org/grpc" 12 | codes "google.golang.org/grpc/codes" 13 | status "google.golang.org/grpc/status" 14 | ) 15 | 16 | // This is a compile-time assertion to ensure that this generated file 17 | // is compatible with the grpc package it is being compiled against. 18 | // Requires gRPC-Go v1.32.0 or later. 19 | const _ = grpc.SupportPackageIsVersion7 20 | 21 | // SleeperClient is the client API for Sleeper service. 22 | // 23 | // For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. 24 | type SleeperClient interface { 25 | Sleep(ctx context.Context, in *SleepRequest, opts ...grpc.CallOption) (*SleepResponse, error) 26 | } 27 | 28 | type sleeperClient struct { 29 | cc grpc.ClientConnInterface 30 | } 31 | 32 | func NewSleeperClient(cc grpc.ClientConnInterface) SleeperClient { 33 | return &sleeperClient{cc} 34 | } 35 | 36 | func (c *sleeperClient) Sleep(ctx context.Context, in *SleepRequest, opts ...grpc.CallOption) (*SleepResponse, error) { 37 | out := new(SleepResponse) 38 | err := c.cc.Invoke(ctx, "/sleepymemory.Sleeper/Sleep", in, out, opts...) 39 | if err != nil { 40 | return nil, err 41 | } 42 | return out, nil 43 | } 44 | 45 | // SleeperServer is the server API for Sleeper service. 46 | // All implementations must embed UnimplementedSleeperServer 47 | // for forward compatibility 48 | type SleeperServer interface { 49 | Sleep(context.Context, *SleepRequest) (*SleepResponse, error) 50 | mustEmbedUnimplementedSleeperServer() 51 | } 52 | 53 | // UnimplementedSleeperServer must be embedded to have forward compatible implementations. 54 | type UnimplementedSleeperServer struct { 55 | } 56 | 57 | func (UnimplementedSleeperServer) Sleep(context.Context, *SleepRequest) (*SleepResponse, error) { 58 | return nil, status.Errorf(codes.Unimplemented, "method Sleep not implemented") 59 | } 60 | func (UnimplementedSleeperServer) mustEmbedUnimplementedSleeperServer() {} 61 | 62 | // UnsafeSleeperServer may be embedded to opt out of forward compatibility for this service. 63 | // Use of this interface is not recommended, as added methods to SleeperServer will 64 | // result in compilation errors. 65 | type UnsafeSleeperServer interface { 66 | mustEmbedUnimplementedSleeperServer() 67 | } 68 | 69 | func RegisterSleeperServer(s grpc.ServiceRegistrar, srv SleeperServer) { 70 | s.RegisterService(&Sleeper_ServiceDesc, srv) 71 | } 72 | 73 | func _Sleeper_Sleep_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { 74 | in := new(SleepRequest) 75 | if err := dec(in); err != nil { 76 | return nil, err 77 | } 78 | if interceptor == nil { 79 | return srv.(SleeperServer).Sleep(ctx, in) 80 | } 81 | info := &grpc.UnaryServerInfo{ 82 | Server: srv, 83 | FullMethod: "/sleepymemory.Sleeper/Sleep", 84 | } 85 | handler := func(ctx context.Context, req interface{}) (interface{}, error) { 86 | return srv.(SleeperServer).Sleep(ctx, req.(*SleepRequest)) 87 | } 88 | return interceptor(ctx, in, info, handler) 89 | } 90 | 91 | // Sleeper_ServiceDesc is the grpc.ServiceDesc for Sleeper service. 92 | // It's only intended for direct use with grpc.RegisterService, 93 | // and not to be introspected or modified (even as a copy) 94 | var Sleeper_ServiceDesc = grpc.ServiceDesc{ 95 | ServiceName: "sleepymemory.Sleeper", 96 | HandlerType: (*SleeperServer)(nil), 97 | Methods: []grpc.MethodDesc{ 98 | { 99 | MethodName: "Sleep", 100 | Handler: _Sleeper_Sleep_Handler, 101 | }, 102 | }, 103 | Streams: []grpc.StreamDesc{}, 104 | Metadata: "sleepymemory/sleepymemory.proto", 105 | } 106 | -------------------------------------------------------------------------------- /sleepyserver/sleepyserver.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "fmt" 7 | "log" 8 | "net" 9 | "net/http" 10 | "net/http/pprof" 11 | "runtime" 12 | "strconv" 13 | "sync" 14 | "time" 15 | 16 | "github.com/evanj/concurrentlimit" 17 | "github.com/evanj/concurrentlimit/sleepymemory" 18 | "golang.org/x/net/netutil" 19 | "google.golang.org/grpc" 20 | "google.golang.org/grpc/codes" 21 | "google.golang.org/grpc/metadata" 22 | "google.golang.org/grpc/status" 23 | "google.golang.org/protobuf/types/known/durationpb" 24 | ) 25 | 26 | const sleepHTTPKey = "sleep" 27 | const wasteHTTPKey = "waste" 28 | 29 | type server struct { 30 | sleepymemory.UnimplementedSleeperServer 31 | logger concurrentMaxLogger 32 | limiter concurrentlimit.Limiter 33 | logAllRequests bool 34 | } 35 | 36 | func newServer(limiter concurrentlimit.Limiter, logAllRequests bool) *server { 37 | return &server{ 38 | logger: concurrentMaxLogger{}, 39 | limiter: limiter, 40 | logAllRequests: logAllRequests, 41 | } 42 | } 43 | 44 | func (s *server) rawRootHandler(w http.ResponseWriter, r *http.Request) { 45 | if r.URL.Path != "/" { 46 | http.NotFound(w, r) 47 | return 48 | } 49 | if r.Method != http.MethodGet { 50 | http.Error(w, "only GET is supported", http.StatusMethodNotAllowed) 51 | return 52 | } 53 | 54 | err := s.rootHandler(w, r) 55 | if err != nil { 56 | statusCode := http.StatusInternalServerError 57 | if err == concurrentlimit.ErrLimited { 58 | statusCode = http.StatusTooManyRequests 59 | } 60 | http.Error(w, err.Error(), statusCode) 61 | } 62 | } 63 | 64 | func humanBytes(bytes uint64) string { 65 | megabytes := float64(bytes) / float64(1024*1024) 66 | return fmt.Sprintf("%.1f", megabytes) 67 | } 68 | 69 | func (s *server) memstatsHandler(w http.ResponseWriter, r *http.Request) { 70 | stats := &runtime.MemStats{} 71 | runtime.ReadMemStats(stats) 72 | 73 | w.Header().Set("Content-Type", "text/plain;charset=utf-8") 74 | fmt.Fprintf(w, "total bytes of memory obtained from the OS Sys=%d %s\n", 75 | stats.Sys, humanBytes(stats.Sys)) 76 | fmt.Fprintf(w, "bytes of allocated heap objects HeapAlloc=%d %s\n", 77 | stats.HeapAlloc, humanBytes(stats.HeapAlloc)) 78 | } 79 | 80 | func (s *server) rootHandler(w http.ResponseWriter, r *http.Request) error { 81 | req := &sleepymemory.SleepRequest{} 82 | 83 | sleepValue := r.FormValue(sleepHTTPKey) 84 | if sleepValue != "" { 85 | var sleepDuration time.Duration 86 | // try to parse as integer seconds first 87 | seconds, err := strconv.Atoi(sleepValue) 88 | if err == nil { 89 | // SUCCESS! 90 | sleepDuration = time.Duration(seconds) * time.Second 91 | } else { 92 | // fall back to parsing duration, and return that error if it fails 93 | sleepDuration, err = time.ParseDuration(sleepValue) 94 | if err != nil { 95 | return err 96 | } 97 | } 98 | req.SleepDuration = durationpb.New(sleepDuration) 99 | } 100 | 101 | wasteValue := r.FormValue(wasteHTTPKey) 102 | if wasteValue != "" { 103 | bytes, err := strconv.Atoi(wasteValue) 104 | if err != nil { 105 | return err 106 | } 107 | req.WasteBytes = int64(bytes) 108 | } 109 | 110 | resp, err := s.sleepImplementation(r.Context(), req) 111 | if err != nil { 112 | return err 113 | } 114 | 115 | w.Header().Set("Content-Type", "text/plain;charset=utf-8") 116 | fmt.Fprintf(w, "slept for %s (pass ?sleep=x)\nwasted %d bytes (pass ?waste=y)\nignored response=%d\n", 117 | req.SleepDuration.String(), req.WasteBytes, resp.Ignored) 118 | return nil 119 | } 120 | 121 | func (s *server) Sleep(ctx context.Context, request *sleepymemory.SleepRequest) (*sleepymemory.SleepResponse, error) { 122 | resp, err := s.sleepImplementation(ctx, request) 123 | if err == concurrentlimit.ErrLimited { 124 | err = status.Error(codes.ResourceExhausted, err.Error()) 125 | } 126 | return resp, err 127 | } 128 | 129 | func (s *server) sleepImplementation(ctx context.Context, request *sleepymemory.SleepRequest) (*sleepymemory.SleepResponse, error) { 130 | // limit concurrent requests 131 | end, err := s.limiter.Start() 132 | if err != nil { 133 | return nil, err 134 | } 135 | defer end() 136 | 137 | defer s.logger.start()() 138 | 139 | if s.logAllRequests { 140 | md, ok := metadata.FromIncomingContext(ctx) 141 | log.Printf("starting Sleep request=%s md=%v ok=%v", request.String(), md, ok) 142 | } 143 | 144 | wasteSlice := make([]byte, request.WasteBytes) 145 | // touch each page in the slice to ensure it is actually allocated 146 | const pageSize = 4096 147 | for i := 0; i < len(wasteSlice); i += pageSize { 148 | wasteSlice[i] = 0xff 149 | } 150 | 151 | var duration time.Duration 152 | if request.SleepDuration != nil { 153 | if err := request.SleepDuration.CheckValid(); err != nil { 154 | return nil, err 155 | } 156 | duration = request.SleepDuration.AsDuration() 157 | } 158 | // TODO: use ctx for cancellation 159 | time.Sleep(duration) 160 | 161 | // read some of the memory and return it so it doesn't get garbage collected 162 | total := 0 163 | for i := 0; i < len(wasteSlice); i += 10 * pageSize { 164 | total += int(wasteSlice[i]) 165 | } 166 | 167 | return &sleepymemory.SleepResponse{Ignored: int64(total)}, nil 168 | } 169 | 170 | type concurrentMaxLogger struct { 171 | mu sync.Mutex 172 | max int 173 | current int 174 | } 175 | 176 | // start records the start of a concurrent request that is terminated when func is called. 177 | func (c *concurrentMaxLogger) start() func() { 178 | c.mu.Lock() 179 | c.current++ 180 | if c.current > c.max { 181 | c.max = c.current 182 | log.Printf("new max requests=%d", c.max) 183 | } 184 | c.mu.Unlock() 185 | 186 | return c.end 187 | } 188 | 189 | func (c *concurrentMaxLogger) end() { 190 | c.mu.Lock() 191 | c.current-- 192 | if c.current < 0 { 193 | panic("bug: mismatched calls to startRequest/endRequest") 194 | } 195 | c.mu.Unlock() 196 | } 197 | 198 | func main() { 199 | httpAddr := flag.String("httpAddr", "localhost:8080", "Address to listen for HTTP requests") 200 | grpcAddr := flag.String("grpcAddr", "localhost:8081", "Address to listen for gRPC requests") 201 | concurrentRequests := flag.Int("concurrentRequests", 0, "Limits the number of concurrent requests") 202 | concurrentConnections := flag.Int("concurrentConnections", 0, "Limits the number of concurrent connections") 203 | grpcConcurrentStreams := flag.Int("grpcConcurrentStreams", 0, "Limits the number of concurrent connections") 204 | logAll := flag.Bool("logAll", false, "Log all requests") 205 | flag.Parse() 206 | 207 | s := newServer(concurrentlimit.NoLimit(), *logAll) 208 | if *concurrentRequests > 0 { 209 | log.Printf("limiting the server to %d concurrent requests", *concurrentRequests) 210 | s.limiter = concurrentlimit.New(*concurrentRequests) 211 | } 212 | 213 | mux := &http.ServeMux{} 214 | mux.HandleFunc("/", s.rawRootHandler) 215 | mux.HandleFunc("/stats", s.memstatsHandler) 216 | 217 | // copied from http/pprof 218 | mux.HandleFunc("/debug/pprof/", pprof.Index) 219 | mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline) 220 | mux.HandleFunc("/debug/pprof/profile", pprof.Profile) 221 | mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol) 222 | mux.HandleFunc("/debug/pprof/trace", pprof.Trace) 223 | 224 | log.Printf("listening for HTTP on http://%s ...", *httpAddr) 225 | httpListener, err := net.Listen("tcp", *httpAddr) 226 | if err != nil { 227 | panic(err) 228 | } 229 | if *concurrentConnections > 0 { 230 | log.Printf("limiting the HTTP server to %d concurrent connections", *concurrentConnections) 231 | httpListener = netutil.LimitListener(httpListener, *concurrentConnections) 232 | } 233 | 234 | go func() { 235 | err := http.Serve(httpListener, mux) 236 | if err != nil { 237 | panic(err) 238 | } 239 | }() 240 | 241 | log.Printf("listening for gRPC on grpcAddr=%s ...", *grpcAddr) 242 | grpcListener, err := net.Listen("tcp", *grpcAddr) 243 | if err != nil { 244 | panic(err) 245 | } 246 | if *concurrentConnections > 0 { 247 | log.Printf("limiting the gRPC server to %d concurrent connections", *concurrentConnections) 248 | grpcListener = netutil.LimitListener(grpcListener, *concurrentConnections) 249 | } 250 | 251 | options := []grpc.ServerOption{} 252 | if *grpcConcurrentStreams > 0 { 253 | log.Printf("setting grpc MaxConcurrentStreams=%d", *grpcConcurrentStreams) 254 | options = append(options, grpc.MaxConcurrentStreams(uint32(*grpcConcurrentStreams))) 255 | } 256 | grpcServer := grpc.NewServer(options...) 257 | sleepymemory.RegisterSleeperServer(grpcServer, s) 258 | err = grpcServer.Serve(grpcListener) 259 | if err != nil { 260 | panic(err) 261 | } 262 | } 263 | --------------------------------------------------------------------------------