func TestNewConnPool(t *testing.T) {
t.Parallel()
var numCallbacks int
afterConnect := func(c *pgx.Conn) error {
numCallbacks++
return nil
}
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: 2, AfterConnect: afterConnect}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatal("Unable to establish connection pool")
}
defer pool.Close()
// It initially connects once
stat := pool.Stat()
if stat.CurrentConnections != 1 {
t.Errorf("Expected 1 connection to be established immediately, but %v were", numCallbacks)
}
// Pool creation returns an error if any AfterConnect callback does
errAfterConnect := errors.New("Some error")
afterConnect = func(c *pgx.Conn) error {
return errAfterConnect
}
config = pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: 2, AfterConnect: afterConnect}
pool, err = pgx.NewConnPool(config)
if err != errAfterConnect {
t.Errorf("Expected errAfterConnect but received unexpected: %v", err)
}
}
func (s *S) SetUpSuite(c *C) {
dbname := "controllertest"
if err := pgtestutils.SetupPostgres(dbname); err != nil {
c.Fatal(err)
}
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
Database: dbname,
},
})
if err != nil {
c.Fatal(err)
}
db := postgres.New(pgxpool, nil)
if err = migrateDB(db); err != nil {
c.Fatal(err)
}
// reconnect with que statements prepared now that schema is migrated
pgxpool, err = pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: "/var/run/postgresql",
Database: dbname,
},
AfterConnect: schema.PrepareStatements,
})
if err != nil {
c.Fatal(err)
}
db = postgres.New(pgxpool, nil)
s.flac = newFakeLogAggregatorClient()
s.cc = tu.NewFakeCluster()
s.hc = handlerConfig{
db: db,
cc: s.cc,
lc: s.flac,
rc: newFakeRouter(),
keys: []string{authKey},
}
handler := appHandler(s.hc)
s.srv = httptest.NewServer(handler)
client, err := controller.NewClient(s.srv.URL, authKey)
c.Assert(err, IsNil)
s.c = client
}
func TestOpenFromConnPool(t *testing.T) {
connConfig := pgx.ConnConfig{
Host: "localhost",
User: "******",
Password: "******",
Database: "pgx_test",
}
config := pgx.ConnPoolConfig{ConnConfig: connConfig}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatalf("Unable to create connection pool: %v", err)
}
defer pool.Close()
db, err := stdlib.OpenFromConnPool(pool)
if err != nil {
t.Fatalf("Unable to create connection pool: %v", err)
}
defer closeDB(t, db)
// Can get pgx.ConnPool from driver
driver := db.Driver().(*stdlib.Driver)
if driver.Pool == nil {
t.Fatal("Expected driver opened through OpenFromConnPool to have Pool, but it did not")
}
// Normal sql/database still works
var n int64
err = db.QueryRow("select 1").Scan(&n)
if err != nil {
t.Fatalf("db.QueryRow unexpectedly failed: %v", err)
}
}
func (s *S) SetUpSuite(c *C) {
s.discoverd, s.cleanup = setup(c)
dbname := "routertest"
if err := pgtestutils.SetupPostgres(dbname); err != nil {
c.Fatal(err)
}
dsn := fmt.Sprintf("dbname=%s", dbname)
db, err := sql.Open("postgres", dsn)
if err != nil {
c.Fatal(err)
}
if err = migrateDB(db); err != nil {
c.Fatal(err)
}
db.Close()
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
Database: dbname,
},
})
if err != nil {
c.Fatal(err)
}
s.pgx = pgxpool
s.pgx.Exec(sqlCreateTruncateTables)
}
func TestPostgresFilesystem(t *testing.T) {
dbname := "blobstoretest"
if err := pgtestutils.SetupPostgres(dbname); err != nil {
t.Fatal(err)
}
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
Database: dbname,
},
})
if err != nil {
t.Fatal(err)
}
db := postgres.New(pgxpool, nil)
defer db.Close()
fs, err := NewPostgresFilesystem(db)
if err != nil {
t.Fatal(err)
}
testList(fs, t)
testDelete(fs, t)
testOffset(fs, t)
testFilesystem(fs, true, t)
}
func main() {
var err error
connPoolConfig := pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: "127.0.0.1",
User: "******",
Password: "******",
Database: "url_shortener",
Logger: log.New("module", "pgx"),
},
MaxConnections: 5,
AfterConnect: afterConnect,
}
pool, err = pgx.NewConnPool(connPoolConfig)
if err != nil {
log.Crit("Unable to create connection pool", "error", err)
os.Exit(1)
}
http.HandleFunc("/", urlHandler)
log.Info("Starting URL shortener on localhost:8080")
err = http.ListenAndServe("localhost:8080", nil)
if err != nil {
log.Crit("Unable to start web server", "error", err)
os.Exit(1)
}
}
func createConnPool(t *testing.T, maxConnections int) *pgx.ConnPool {
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: maxConnections}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatalf("Unable to create connection pool: %v", err)
}
return pool
}
func openTestClientMaxConns(t testing.TB, maxConnections int) *Client {
connPoolConfig := pgx.ConnPoolConfig{
ConnConfig: testConnConfig,
MaxConnections: maxConnections,
AfterConnect: PrepareStatements,
}
pool, err := pgx.NewConnPool(connPoolConfig)
if err != nil {
t.Fatal(err)
}
return NewClient(pool)
}
func TestNewConnPoolMaxConnectionsCannotBeLessThan2(t *testing.T) {
t.Parallel()
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: 1}
pool, err := pgx.NewConnPool(config)
if err == nil {
pool.Close()
t.Fatal(`Expected NewConnPool to fail with "MaxConnections must be at least 2" error, but it succeeded`)
}
if err.Error() != "MaxConnections must be at least 2" {
t.Fatalf(`Expected NewConnPool to fail with "MaxConnections must be at least 2" error, but it failed with %v`, err)
}
}
func TestNewConnPoolDefaultsTo5MaxConnections(t *testing.T) {
t.Parallel()
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig}
pool, err := pgx.NewConnPool(config)
if err != nil {
t.Fatal("Unable to establish connection pool")
}
defer pool.Close()
if n := pool.Stat().MaxConnections; n != 5 {
t.Fatalf("Expected pool to default to 5 max connections, but it was %d", n)
}
}
func setupTestDB(c *C, dbname string) *postgres.DB {
if err := pgtestutils.SetupPostgres(dbname); err != nil {
c.Fatal(err)
}
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
Database: dbname,
},
})
if err != nil {
c.Fatal(err)
}
return postgres.New(pgxpool, nil)
}
func Open(conf *Conf, afterConn func(*pgx.Conn) error) (*DB, error) {
connConfig := pgx.ConnConfig{
Host: fmt.Sprintf("leader.%s.discoverd", conf.Service),
User: conf.User,
Database: conf.Database,
Password: conf.Password,
Dial: dialer.Retry.Dial,
}
connPool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: connConfig,
AfterConnect: afterConn,
MaxConnections: 20,
})
db := &DB{connPool, conf}
return db, err
}
func BenchmarkConnPool(b *testing.B) {
config := pgx.ConnPoolConfig{ConnConfig: *defaultConnConfig, MaxConnections: 5}
pool, err := pgx.NewConnPool(config)
if err != nil {
b.Fatalf("Unable to create connection pool: %v", err)
}
defer pool.Close()
b.ResetTimer()
for i := 0; i < b.N; i++ {
var conn *pgx.Conn
if conn, err = pool.Acquire(); err != nil {
b.Fatalf("Unable to acquire connection: %v", err)
}
pool.Release(conn)
}
}
func main() {
log := logger.New("fn", "main")
log.Info("creating controller client")
client, err := controller.NewClient("", os.Getenv("AUTH_KEY"))
if err != nil {
log.Error("error creating controller client", "err", err)
shutdown.Fatal()
}
log.Info("connecting to postgres")
db := postgres.Wait("", "")
log.Info("creating postgres connection pool")
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
User: os.Getenv("PGUSER"),
Password: os.Getenv("PGPASSWORD"),
Database: os.Getenv("PGDATABASE"),
},
AfterConnect: que.PrepareStatements,
MaxConnections: workerCount,
})
if err != nil {
log.Error("error creating postgres connection pool", "err", err)
shutdown.Fatal()
}
shutdown.BeforeExit(func() { pgxpool.Close() })
workers := que.NewWorkerPool(
que.NewClient(pgxpool),
que.WorkMap{
"deployment": deployment.JobHandler(db, client, logger),
"app_deletion": app_deletion.JobHandler(db, client, logger),
},
workerCount,
)
workers.Interval = 5 * time.Second
log.Info("starting workers", "count", workerCount, "interval", workers.Interval)
workers.Start()
shutdown.BeforeExit(func() { workers.Shutdown() })
select {} // block and keep running
}
func (s *S) SetUpSuite(c *C) {
s.discoverd, s.cleanup = setup(c)
if err := pgtestutils.SetupPostgres(dbname); err != nil {
c.Fatal(err)
}
pgxpool, err := pgx.NewConnPool(newPgxConnPoolConfig())
if err != nil {
c.Fatal(err)
}
db := postgres.New(pgxpool, nil)
if err = migrateDB(db); err != nil {
c.Fatal(err)
}
s.pgx = db.ConnPool
s.pgx.Exec(sqlCreateTruncateTables)
}
func (s *S) SetUpSuite(c *C) {
dbname := "controllertest"
if err := pgtestutils.SetupPostgres(dbname); err != nil {
c.Fatal(err)
}
dsn := fmt.Sprintf("dbname=%s", dbname)
db, err := sql.Open("postgres", dsn)
if err != nil {
c.Fatal(err)
}
if err = migrateDB(db); err != nil {
c.Fatal(err)
}
pg := postgres.New(db, dsn)
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: "/var/run/postgresql",
Database: dbname,
},
AfterConnect: que.PrepareStatements,
})
if err != nil {
c.Fatal(err)
}
s.flac = newFakeLogAggregatorClient()
s.cc = tu.NewFakeCluster()
s.hc = handlerConfig{
db: pg,
cc: s.cc,
lc: s.flac,
rc: newFakeRouter(),
pgxpool: pgxpool,
keys: []string{authKey},
}
handler := appHandler(s.hc)
s.srv = httptest.NewServer(handler)
client, err := controller.NewClient(s.srv.URL, authKey)
c.Assert(err, IsNil)
s.c = client
}
func (s *S) SetUpSuite(c *C) {
dbname := "controllertest"
db := setupTestDB(c, dbname)
if err := migrateDB(db); err != nil {
c.Fatal(err)
}
// reconnect with que statements prepared now that schema is migrated
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: "/var/run/postgresql",
Database: dbname,
},
AfterConnect: schema.PrepareStatements,
})
if err != nil {
c.Fatal(err)
}
db = postgres.New(pgxpool, nil)
ca, err := certgen.Generate(certgen.Params{IsCA: true})
if err != nil {
c.Fatal(err)
}
s.caCert = []byte(ca.PEM)
s.flac = newFakeLogAggregatorClient()
s.cc = tu.NewFakeCluster()
s.hc = handlerConfig{
db: db,
cc: s.cc,
lc: s.flac,
rc: newFakeRouter(),
keys: []string{authKey},
caCert: s.caCert,
}
handler := appHandler(s.hc)
s.srv = httptest.NewServer(handler)
client, err := controller.NewClient(s.srv.URL, authKey)
c.Assert(err, IsNil)
s.c = client
}
func main() {
var err error
pool, err = pgx.NewConnPool(extractConfig())
if err != nil {
fmt.Fprintln(os.Stderr, "Unable to connect to database:", err)
os.Exit(1)
}
go listen()
fmt.Println(`Type a message and press enter.
This message should appear in any other chat instances connected to the same
database.
Type "exit" to quit.
`)
scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
msg := scanner.Text()
if msg == "exit" {
os.Exit(0)
}
_, err = pool.Exec("select pg_notify('chat', $1)", msg)
if err != nil {
fmt.Fprintln(os.Stderr, "Error sending notification:", err)
os.Exit(1)
}
}
if err := scanner.Err(); err != nil {
fmt.Fprintln(os.Stderr, "Error scanning from stdin:", err)
os.Exit(1)
}
}
func NewPostgresFilesystem(db *sql.DB) (Filesystem, error) {
m := postgres.NewMigrations()
m.Add(1,
`CREATE TABLE files (
file_id oid PRIMARY KEY DEFAULT lo_create(0),
name text UNIQUE NOT NULL,
size bigint,
type text,
digest text,
created_at timestamp with time zone NOT NULL DEFAULT current_timestamp
);`,
`CREATE FUNCTION delete_file() RETURNS TRIGGER AS $$
BEGIN
PERFORM lo_unlink(OLD.file_id);
RETURN NULL;
END;
$$ LANGUAGE plpgsql;`,
`CREATE TRIGGER delete_file
AFTER DELETE ON files
FOR EACH ROW EXECUTE PROCEDURE delete_file();`,
)
// TODO(jpg) reuse pkg/postgres connection when converted
connConf := pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
User: os.Getenv("PGUSER"),
Password: os.Getenv("PGPASSWORD"),
Database: os.Getenv("PGDATABASE"),
}
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: connConf,
})
if err != nil {
return nil, err
}
return &PostgresFilesystem{db: pgxpool}, m.Migrate(db)
}
func (s *S) SetUpSuite(c *C) {
s.discoverd, s.cleanup = setup(c)
dbname := "routertest"
if err := pgtestutils.SetupPostgres(dbname); err != nil {
c.Fatal(err)
}
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
Database: dbname,
},
})
if err != nil {
c.Fatal(err)
}
db := postgres.New(pgxpool, nil)
if err = migrateDB(db); err != nil {
c.Fatal(err)
}
s.pgx = db.ConnPool
s.pgx.Exec(sqlCreateTruncateTables)
}
func main() {
defer shutdown.Exit()
port := os.Getenv("PORT")
if port == "" {
port = "3000"
}
addr := ":" + port
if seed := os.Getenv("NAME_SEED"); seed != "" {
s, err := hex.DecodeString(seed)
if err != nil {
log.Fatalln("error decoding NAME_SEED:", err)
}
name.SetSeed(s)
}
db := postgres.Wait("", "")
if err := migrateDB(db.DB); err != nil {
shutdown.Fatal(err)
}
pgxcfg, err := pgx.ParseURI(fmt.Sprintf("http://%s:%s@%s/%s", os.Getenv("PGUSER"), os.Getenv("PGPASSWORD"), db.Addr(), os.Getenv("PGDATABASE")))
if err != nil {
log.Fatal(err)
}
pgxcfg.Dial = dialer.Retry.Dial
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgxcfg,
AfterConnect: que.PrepareStatements,
})
if err != nil {
log.Fatal(err)
}
shutdown.BeforeExit(func() { pgxpool.Close() })
lc, err := logaggc.New("")
if err != nil {
shutdown.Fatal(err)
}
rc := routerc.New()
hb, err := discoverd.DefaultClient.AddServiceAndRegisterInstance("flynn-controller", &discoverd.Instance{
Addr: addr,
Proto: "http",
Meta: map[string]string{
"AUTH_KEY": os.Getenv("AUTH_KEY"),
},
})
if err != nil {
shutdown.Fatal(err)
}
shutdown.BeforeExit(func() {
hb.Close()
})
handler := appHandler(handlerConfig{
db: db,
cc: clusterClientWrapper{cluster.NewClient()},
lc: lc,
rc: rc,
pgxpool: pgxpool,
keys: strings.Split(os.Getenv("AUTH_KEY"), ","),
})
shutdown.Fatal(http.ListenAndServe(addr, handler))
}
func main() {
log := logger.New("fn", "main")
log.Info("creating controller client")
client, err := controller.NewClient("", os.Getenv("AUTH_KEY"))
if err != nil {
log.Error("error creating controller client", "err", err)
shutdown.Fatal()
}
log.Info("connecting to postgres")
db := postgres.Wait("", "")
log.Info("creating postgres connection pool")
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
User: os.Getenv("PGUSER"),
Password: os.Getenv("PGPASSWORD"),
Database: os.Getenv("PGDATABASE"),
},
AfterConnect: que.PrepareStatements,
MaxConnections: workerCount,
})
if err != nil {
log.Error("error creating postgres connection pool", "err", err)
shutdown.Fatal()
}
shutdown.BeforeExit(func() { pgxpool.Close() })
go func() {
status.AddHandler(func() status.Status {
_, err := pgxpool.Exec("SELECT 1")
if err != nil {
return status.Unhealthy
}
return status.Healthy
})
addr := ":" + os.Getenv("PORT")
hb, err := discoverd.AddServiceAndRegister("flynn-controller-worker", addr)
if err != nil {
shutdown.Fatal(err)
}
shutdown.BeforeExit(func() { hb.Close() })
shutdown.Fatal(http.ListenAndServe(addr, nil))
}()
workers := que.NewWorkerPool(
que.NewClient(pgxpool),
que.WorkMap{
"deployment": deployment.JobHandler(db, client, logger),
"app_deletion": app_deletion.JobHandler(db, client, logger),
},
workerCount,
)
workers.Interval = 5 * time.Second
log.Info("starting workers", "count", workerCount, "interval", workers.Interval)
workers.Start()
shutdown.BeforeExit(func() { workers.Shutdown() })
select {} // block and keep running
}
func main() {
defer shutdown.Exit()
var cookieKey *[32]byte
if key := os.Getenv("COOKIE_KEY"); key != "" {
res, err := base64.StdEncoding.DecodeString(key)
if err != nil {
shutdown.Fatalf("error decoding COOKIE_KEY: %s", err)
}
if len(res) != 32 {
shutdown.Fatalf("decoded %d bytes from COOKIE_KEY, expected 32", len(res))
}
var k [32]byte
copy(k[:], res)
cookieKey = &k
}
if cookieKey == nil {
shutdown.Fatal("Missing random 32 byte base64-encoded COOKIE_KEY")
}
httpPort := flag.String("http-port", "8080", "http listen port")
httpsPort := flag.String("https-port", "4433", "https listen port")
tcpIP := flag.String("tcp-ip", os.Getenv("LISTEN_IP"), "tcp router listen ip")
tcpRangeStart := flag.Int("tcp-range-start", 3000, "tcp port range start")
tcpRangeEnd := flag.Int("tcp-range-end", 3500, "tcp port range end")
certFile := flag.String("tls-cert", "", "TLS (SSL) cert file in pem format")
keyFile := flag.String("tls-key", "", "TLS (SSL) key file in pem format")
apiPort := flag.String("api-port", "", "api listen port")
flag.Parse()
if *apiPort == "" {
*apiPort = os.Getenv("PORT")
if *apiPort == "" {
*apiPort = "5000"
}
}
keypair := tls.Certificate{}
var err error
if *certFile != "" {
if keypair, err = tls.LoadX509KeyPair(*certFile, *keyFile); err != nil {
shutdown.Fatal(err)
}
} else if tlsCert := os.Getenv("TLSCERT"); tlsCert != "" {
if tlsKey := os.Getenv("TLSKEY"); tlsKey != "" {
os.Setenv("TLSKEY", fmt.Sprintf("md5^(%s)", md5sum(tlsKey)))
if keypair, err = tls.X509KeyPair([]byte(tlsCert), []byte(tlsKey)); err != nil {
shutdown.Fatal(err)
}
}
}
log := logger.New("fn", "main")
log.Info("connecting to postgres")
db, err := postgres.Open("", "")
if err != nil {
log.Error("error connecting to postgres", "err", err)
shutdown.Fatal(err)
}
log.Info("running DB migrations")
if err := migrateDB(db.DB); err != nil {
log.Error("error running DB migrations", "err", err)
shutdown.Fatal(err)
}
var pgport int
if port := os.Getenv("PGPORT"); port != "" {
var err error
if pgport, err = strconv.Atoi(port); err != nil {
shutdown.Fatal(err)
}
}
log.Info("creating postgres connection pool")
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
Port: uint16(pgport),
Database: os.Getenv("PGDATABASE"),
User: os.Getenv("PGUSER"),
Password: os.Getenv("PGPASSWORD"),
},
})
if err != nil {
log.Error("error creating postgres connection pool", "err", err)
shutdown.Fatal(err)
}
shutdown.BeforeExit(func() { pgxpool.Close() })
httpAddr := net.JoinHostPort(os.Getenv("LISTEN_IP"), *httpPort)
httpsAddr := net.JoinHostPort(os.Getenv("LISTEN_IP"), *httpsPort)
r := Router{
TCP: &TCPListener{
IP: *tcpIP,
startPort: *tcpRangeStart,
endPort: *tcpRangeEnd,
ds: NewPostgresDataStore("tcp", pgxpool),
discoverd: discoverd.DefaultClient,
},
HTTP: &HTTPListener{
Addr: httpAddr,
TLSAddr: httpsAddr,
cookieKey: cookieKey,
keypair: keypair,
ds: NewPostgresDataStore("http", pgxpool),
discoverd: discoverd.DefaultClient,
},
}
if err := r.Start(); err != nil {
shutdown.Fatal(err)
}
shutdown.BeforeExit(r.Close)
apiAddr := net.JoinHostPort(os.Getenv("LISTEN_IP"), *apiPort)
log.Info("starting API listener")
listener, err := listenFunc("tcp4", apiAddr)
if err != nil {
log.Error("error starting API listener", "err", err)
shutdown.Fatal(listenErr{apiAddr, err})
}
services := map[string]string{
"router-api": apiAddr,
"router-http": httpAddr,
}
for service, addr := range services {
log.Info("registering service", "name", service, "addr", addr)
hb, err := discoverd.AddServiceAndRegister(service, addr)
if err != nil {
log.Error("error registering service", "name", service, "addr", addr, "err", err)
shutdown.Fatal(err)
}
shutdown.BeforeExit(func() { hb.Close() })
}
log.Info("serving API requests")
shutdown.Fatal(http.Serve(listener, apiHandler(&r)))
}
func main() {
defer shutdown.Exit()
apiPort := os.Getenv("PORT")
if apiPort == "" {
apiPort = "5000"
}
var cookieKey *[32]byte
if key := os.Getenv("COOKIE_KEY"); key != "" {
res, err := base64.StdEncoding.DecodeString(key)
if err != nil {
shutdown.Fatal("error decoding COOKIE_KEY:", err)
}
var k [32]byte
copy(k[:], res)
cookieKey = &k
}
httpAddr := flag.String("httpaddr", ":8080", "http listen address")
httpsAddr := flag.String("httpsaddr", ":4433", "https listen address")
tcpIP := flag.String("tcpip", "", "tcp router listen ip")
tcpRangeStart := flag.Int("tcp-range-start", 3000, "tcp port range start")
tcpRangeEnd := flag.Int("tcp-range-end", 3500, "tcp port range end")
certFile := flag.String("tlscert", "", "TLS (SSL) cert file in pem format")
keyFile := flag.String("tlskey", "", "TLS (SSL) key file in pem format")
apiAddr := flag.String("apiaddr", ":"+apiPort, "api listen address")
flag.Parse()
keypair := tls.Certificate{}
var err error
if *certFile != "" {
if keypair, err = tls.LoadX509KeyPair(*certFile, *keyFile); err != nil {
shutdown.Fatal(err)
}
} else if tlsCert := os.Getenv("TLSCERT"); tlsCert != "" {
if tlsKey := os.Getenv("TLSKEY"); tlsKey != "" {
os.Setenv("TLSKEY", fmt.Sprintf("md5^(%s)", md5sum(tlsKey)))
if keypair, err = tls.X509KeyPair([]byte(tlsCert), []byte(tlsKey)); err != nil {
shutdown.Fatal(err)
}
}
}
db, err := postgres.Open("", "")
if err != nil {
shutdown.Fatal(err)
}
if err := migrateDB(db.DB); err != nil {
shutdown.Fatal(err)
}
var pgport int
if port := os.Getenv("PGPORT"); port != "" {
var err error
if pgport, err = strconv.Atoi(port); err != nil {
shutdown.Fatal(err)
}
}
pgxpool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: os.Getenv("PGHOST"),
Port: uint16(pgport),
Database: os.Getenv("PGDATABASE"),
User: os.Getenv("PGUSER"),
Password: os.Getenv("PGPASSWORD"),
},
})
if err != nil {
shutdown.Fatal(err)
}
shutdown.BeforeExit(func() { pgxpool.Close() })
r := Router{
TCP: &TCPListener{
IP: *tcpIP,
startPort: *tcpRangeStart,
endPort: *tcpRangeEnd,
ds: NewPostgresDataStore("tcp", pgxpool),
discoverd: discoverd.DefaultClient,
},
HTTP: &HTTPListener{
Addr: *httpAddr,
TLSAddr: *httpsAddr,
cookieKey: cookieKey,
keypair: keypair,
ds: NewPostgresDataStore("http", pgxpool),
discoverd: discoverd.DefaultClient,
},
}
if err := r.Start(); err != nil {
shutdown.Fatal(err)
}
listener, err := listenFunc("tcp4", *apiAddr)
if err != nil {
shutdown.Fatal(listenErr{*apiAddr, err})
}
services := map[string]string{
"router-api": *apiAddr,
"router-http": *httpAddr,
}
for service, addr := range services {
hb, err := discoverd.AddServiceAndRegister(service, addr)
if err != nil {
shutdown.Fatal(err)
}
shutdown.BeforeExit(func() { hb.Close() })
}
shutdown.Fatal(http.Serve(listener, apiHandler(&r)))
}
func (s *S) TestHTTPResync(c *C) {
var connPids []int32
var cmu sync.Mutex
poolConfig := newPgxConnPoolConfig()
poolConfig.AfterConnect = func(conn *pgx.Conn) error {
cmu.Lock()
defer cmu.Unlock()
connPids = append(connPids, conn.Pid)
return nil
}
pgxpool, err := pgx.NewConnPool(poolConfig)
if err != nil {
c.Fatal(err)
}
l := &HTTPListener{
Addr: "127.0.0.1:0",
ds: NewPostgresDataStore("http", pgxpool),
discoverd: s.discoverd,
}
if err := l.Start(); err != nil {
c.Fatal(err)
}
defer l.Close()
srv := httptest.NewServer(httpTestHandler("1"))
srv2 := httptest.NewServer(httpTestHandler("2"))
defer srv.Close()
defer srv2.Close()
route := addRoute(c, l, router.HTTPRoute{
Domain: "example.com",
Service: "example-com",
}.ToRoute())
discoverdRegisterHTTPService(c, l, "example-com", srv.Listener.Addr().String())
assertGet(c, "http://"+l.Addr, "example.com", "1")
// testing hooks
presyncc := make(chan struct{})
l.preSync = func() {
<-presyncc
}
postsyncc := make(chan struct{})
l.postSync = func(startc <-chan struct{}) {
<-startc
close(postsyncc)
}
terminateAllPids := func() {
cmu.Lock()
defer cmu.Unlock()
// grab a fresh conn
conn, err := pgx.Connect(poolConfig.ConnConfig)
if err != nil {
c.Fatal(err)
}
defer conn.Close()
// terminate all conns from the pool
for _, pid := range connPids {
if _, err := conn.Exec("select pg_terminate_backend($1)", pid); err != nil {
c.Fatalf("Unable to kill backend PostgreSQL process: %v", err)
}
}
connPids = connPids[:0]
}
terminateAllPids()
// Because we terminated all of the connections some will be dead, others
// will become dead when we try to write to them.
attempts := 0
for {
if attempts > 5 {
c.Fatal(fmt.Errorf("Unable to remove route after disconnecting sync"))
}
err = l.RemoveRoute(route.ID)
if e, ok := err.(*net.OpError); ok {
if ee, ok := e.Err.(*os.SyscallError); ok && ee.Err == syscall.EPIPE || e.Err == syscall.EPIPE {
attempts++
continue
}
}
if err == pgx.ErrDeadConn {
attempts++
continue
}
c.Assert(err, IsNil)
break
}
// Also add a new route
err = l.AddRoute(router.HTTPRoute{
Domain: "example.org",
Service: "example-org",
}.ToRoute())
c.Assert(err, IsNil)
// trigger the reconnect
close(presyncc)
// wait for the sync to complete
<-postsyncc
// ensure that route was actually removed
res, err := httpClient.Do(newReq("http://"+l.Addr, "example.com"))
c.Assert(err, IsNil)
c.Assert(res.StatusCode, Equals, 404)
res.Body.Close()
// ensure that new route was added and traffic being directed
discoverdRegisterHTTPService(c, l, "example-org", srv2.Listener.Addr().String())
assertGet(c, "http://"+l.Addr, "example.org", "2")
}
func (p *Postgres) start() error {
log := p.log.New("fn", "start", "data_dir", p.dataDir, "bin_dir", p.binDir)
log.Info("starting postgres")
// clear stale pid if it exists
os.Remove(filepath.Join(p.dataDir, "postmaster.pid"))
p.expectExit.Store(false)
p.daemonExit = make(chan struct{})
cmd := exec.Command(p.binPath("postgres"), "-D", p.dataDir)
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
if err := cmd.Start(); err != nil {
log.Error("failed to start postgres", "err", err)
return err
}
p.daemon = cmd
p.setRunning(true)
go func() {
err := cmd.Wait()
if !p.expectExit.Load().(bool) {
p.log.Error("postgres unexpectedly exit", "err", err)
shutdown.ExitWithCode(1)
}
close(p.daemonExit)
}()
log.Debug("waiting for postgres to start")
startTime := time.Now().UTC()
var err error
for {
if time.Now().Sub(startTime) > p.opTimeout {
log.Error("timed out waiting for postgres to start", "err", err)
if err := p.stop(); err != nil {
log.Error("error stopping postgres", "err", err)
}
return err
}
port, _ := strconv.Atoi(p.port)
c := pgx.ConnPoolConfig{
ConnConfig: pgx.ConnConfig{
Host: "127.0.0.1",
User: "******",
Port: uint16(port),
},
}
p.dbMtx.Lock()
p.db, err = pgx.NewConnPool(c)
p.dbMtx.Unlock()
if err == nil {
_, err = p.db.Exec("SELECT 1")
if err == nil {
log.Info("postgres started")
return nil
}
}
log.Debug("ignoring error connecting to postgres", "err", err)
time.Sleep(checkInterval)
}
}