regrace/gracehttp/testbin_test.go at master · coderme/regrace · GitHub

History

226 lines (209 loc) · 6.58 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

package gracehttp_test

import (

"crypto/tls"

"encoding/json"

"flag"

"fmt"

"log"

"net/http"

"os"

"strings"

"sync"

"testing"

"time"

"github.com/facebookgo/grace/gracehttp"

)

const preStartProcessEnv = "GRACEHTTP_PRE_START_PROCESS"

func TestMain(m *testing.M) {

const (

testbinKey = "GRACEHTTP_TEST_BIN"

testbinValue = "1"

)

if os.Getenv(testbinKey) == testbinValue {

testbinMain()

return

}

if err := os.Setenv(testbinKey, testbinValue); err != nil {

panic(err)

}

os.Exit(m.Run())

}

type response struct {

Sleep time.Duration

Pid int

Error string `json:",omitempty"`

}

// Wait for 10 consecutive responses from our own pid.

// This prevents flaky tests that arise from the fact that we have the

// perfectly acceptable (read: not a bug) condition where both the new and the

// old servers are accepting requests. In fact the amount of time both are

// accepting at the same time and the number of requests that flip flop between

// them is unbounded and in the hands of the various kernels our code tends to

// run on.

// In order to combat this, we wait for 10 successful responses from our own

// pid. This is a somewhat reliable way to ensure the old server isn't

// serving anymore.

func wait(wg *sync.WaitGroup, url string) {

var success int

defer wg.Done()

for {

res, err := http.Get(url)

if err == nil {

// ensure it isn't a response from a previous instance

defer res.Body.Close()

var r response

if err := json.NewDecoder(res.Body).Decode(&r); err != nil {

log.Fatalf("Error decoding json: %s", err)

}

if r.Pid == os.Getpid() {

success++

if success == 10 {

return

}

continue

}

} else {

success = 0

// we expect connection refused

if !strings.HasSuffix(err.Error(), "connection refused") {

e2 := json.NewEncoder(os.Stderr).Encode(&response{

Error: err.Error(),

Pid: os.Getpid(),

})

if e2 != nil {

log.Fatalf("Error writing error json: %s", e2)

}

func httpsServer(addr string) *http.Server {

cert, err := tls.X509KeyPair(localhostCert, localhostKey)

if err != nil {

log.Fatalf("error loading cert: %v", err)

}

return &http.Server{

Addr: addr,

Handler: newHandler(),

TLSConfig: &tls.Config{

NextProtos: []string{"http/1.1"},

Certificates: []tls.Certificate{cert},

}

func testbinMain() {

var httpAddr, httpsAddr string

var testOption int

flag.StringVar(&httpAddr, "http", ":48560", "http address to bind to")

flag.StringVar(&httpsAddr, "https", ":48561", "https address to bind to")

flag.IntVar(&testOption, "testOption", -1, "which option to test on ServeWithOptions")

flag.Parse()

// we have self signed certs

http.DefaultTransport = &http.Transport{

DisableKeepAlives: true,

TLSClientConfig: &tls.Config{

InsecureSkipVerify: true,

}

// print json to stderr once we can successfully connect to all three

// addresses. the ensures we only print the line once we're ready to serve.

go func() {

var wg sync.WaitGroup

wg.Add(2)

go wait(&wg, fmt.Sprintf("http://%s/sleep/?duration=1ms", httpAddr))

go wait(&wg, fmt.Sprintf("https://%s/sleep/?duration=1ms", httpsAddr))

wg.Wait()

err := json.NewEncoder(os.Stderr).Encode(&response{Pid: os.Getpid()})

if err != nil {

log.Fatalf("Error writing startup json: %s", err)

}

}()

servers := []*http.Server{

&http.Server{Addr: httpAddr, Handler: newHandler()},

httpsServer(httpsAddr),

}

if testOption == -1 {

err := gracehttp.Serve(servers...)

if err != nil {

log.Fatalf("Error in gracehttp.Serve: %s", err)

}

} else {

if testOption == testPreStartProcess {

switch os.Getenv(preStartProcessEnv) {

case "":

err := os.Setenv(preStartProcessEnv, "READY")

if err != nil {

log.Fatalf("testbin (first incarnation) could not set %v to 'ready': %v", preStartProcessEnv, err)

}

case "FIRED":

// all good, reset for next round

err := os.Setenv(preStartProcessEnv, "READY")

if err != nil {

log.Fatalf("testbin (second incarnation) could not reset %v to 'ready': %v", preStartProcessEnv, err)

}

case "READY":

log.Fatalf("failure to update startup hook before new process started")

default:

log.Fatalf("something strange happened with %v: it ended up as %v, which is not '', 'FIRED', or 'READY'", preStartProcessEnv, os.Getenv(preStartProcessEnv))

}

err := gracehttp.ServeWithOptions(

servers,

gracehttp.PreStartProcess(func() error {

err := os.Setenv(preStartProcessEnv, "FIRED")

if err != nil {

log.Fatalf("startup hook could not set %v to 'fired': %v", preStartProcessEnv, err)

}

return nil

}),

)

if err != nil {

log.Fatalf("Error in gracehttp.Serve: %s", err)

}

func newHandler() http.Handler {

mux := http.NewServeMux()

mux.HandleFunc("/sleep/", func(w http.ResponseWriter, r *http.Request) {

duration, err := time.ParseDuration(r.FormValue("duration"))

if err != nil {

http.Error(w, err.Error(), 400)

}

time.Sleep(duration)

err = json.NewEncoder(w).Encode(&response{

Sleep: duration,

Pid: os.Getpid(),

})

if err != nil {

log.Fatalf("Error encoding json: %s", err)

}

})

return mux

}

// localhostCert is a PEM-encoded TLS cert with SAN IPs

// "127.0.0.1" and "[::1]", expiring at the last second of 2049 (the end

// of ASN.1 time).

// generated from src/pkg/crypto/tls:

// go run generate_cert.go --rsa-bits 512 --host 127.0.0.1,::1,example.com --ca --start-date "Jan 1 00:00:00 1970" --duration=1000000h

var localhostCert = []byte(`-----BEGIN CERTIFICATE-----

MIIBdzCCASOgAwIBAgIBADALBgkqhkiG9w0BAQUwEjEQMA4GA1UEChMHQWNtZSBD

bzAeFw03MDAxMDEwMDAwMDBaFw00OTEyMzEyMzU5NTlaMBIxEDAOBgNVBAoTB0Fj

bWUgQ28wWjALBgkqhkiG9w0BAQEDSwAwSAJBALyCfqwwip8BvTKgVKGdmjZTU8DD

ndR+WALmFPIRqn89bOU3s30olKiqYEju/SFoEvMyFRT/TWEhXHDaufThqaMCAwEA

AaNoMGYwDgYDVR0PAQH/BAQDAgCkMBMGA1UdJQQMMAoGCCsGAQUFBwMBMA8GA1Ud

EwEB/wQFMAMBAf8wLgYDVR0RBCcwJYILZXhhbXBsZS5jb22HBH8AAAGHEAAAAAAA

AAAAAAAAAAAAAAEwCwYJKoZIhvcNAQEFA0EAr/09uy108p51rheIOSnz4zgduyTl

M+4AmRo8/U1twEZLgfAGG/GZjREv2y4mCEUIM3HebCAqlA5jpRg76Rf8jw==

-----END CERTIFICATE-----`)

// localhostKey is the private key for localhostCert.

var localhostKey = []byte(`-----BEGIN RSA PRIVATE KEY-----

MIIBOQIBAAJBALyCfqwwip8BvTKgVKGdmjZTU8DDndR+WALmFPIRqn89bOU3s30o

lKiqYEju/SFoEvMyFRT/TWEhXHDaufThqaMCAwEAAQJAPXuWUxTV8XyAt8VhNQER

LgzJcUKb9JVsoS1nwXgPksXnPDKnL9ax8VERrdNr+nZbj2Q9cDSXBUovfdtehcdP

qQIhAO48ZsPylbTrmtjDEKiHT2Ik04rLotZYS2U873J6I7WlAiEAypDjYxXyafv/

Yo1pm9onwcetQKMW8CS3AjuV9Axzj6cCIEx2Il19fEMG4zny0WPlmbrcKvD/DpJQ

4FHrzsYlIVTpAiAas7S1uAvneqd0l02HlN9OxQKKlbUNXNme+rnOnOGS2wIgS0jW

zl1jvrOSJeP1PpAHohWz6LOhEr8uvltWkN6x3vE=

-----END RSA PRIVATE KEY-----`)

regrace/gracehttp/testbin_test.go at master · coderme/regrace · GitHub

Expand file tree

Breadcrumbs

testbin_test.go

Latest commit

History

Breadcrumbs

testbin_test.go

File metadata and controls