Example #1
0
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)
	}
}
Example #2
0
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
}
Example #3
0
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)
	}
}
Example #4
0
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)
}
Example #5
0
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)
}
Example #6
0
File: main.go Project: devick/flynn
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)
	}
}
Example #7
0
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
}
Example #8
0
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)
}
Example #9
0
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)
	}
}
Example #10
0
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)
	}
}
Example #11
0
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)
}
Example #12
0
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
}
Example #13
0
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)
	}
}
Example #14
0
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
}
Example #15
0
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)
}
Example #16
0
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
}
Example #17
0
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
}
Example #18
0
File: main.go Project: devick/flynn
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)
	}
}
Example #19
0
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)
}
Example #20
0
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)
}
Example #21
0
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))
}
Example #22
0
File: main.go Project: ozum/flynn
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
}
Example #23
0
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)))
}
Example #24
0
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)))
}
Example #25
0
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")
}
Example #26
0
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)
	}
}