Exemplo n.º 1
0
func (s *sphinxPools) newConn() (*SphinxDB, error) {
	var (
		err error
		db  *sql.DB
		o   orm.Ormer
	)

	if db, err = sql.Open("sphinx", "root:root@tcp("+setting.SphinxHost+")/?loc=UTC"); err != nil {
		return nil, err
	}

	if err = db.Ping(); err != nil {
		db.Close()
		return nil, err
	}

	db.SetMaxIdleConns(1)
	db.SetMaxOpenConns(2)

	o, err = orm.NewOrmWithDB("sphinx", "sphinx", db)
	if err != nil {
		return nil, err
	}

	sdb := &SphinxDB{
		alive: true,
		pools: s,
		db:    db,
		orm:   o,
	}
	return sdb, nil
}
Exemplo n.º 2
0
func postgresSetup(db *sql.DB) error {
	var err error

	db.SetMaxIdleConns(4)
	db.SetMaxOpenConns(8)

	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS Devices (
		DeviceID CHAR(63) NOT NULL PRIMARY KEY,
		Seen TIMESTAMP NOT NULL
	)`)
	if err != nil {
		return err
	}

	var tmp string
	row := db.QueryRow(`SELECT 'DevicesDeviceIDIndex'::regclass`)
	if err = row.Scan(&tmp); err != nil {
		_, err = db.Exec(`CREATE INDEX DevicesDeviceIDIndex ON Devices (DeviceID)`)
	}
	if err != nil {
		return err
	}

	row = db.QueryRow(`SELECT 'DevicesSeenIndex'::regclass`)
	if err = row.Scan(&tmp); err != nil {
		_, err = db.Exec(`CREATE INDEX DevicesSeenIndex ON Devices (Seen)`)
	}
	if err != nil {
		return err
	}

	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS Addresses (
		DeviceID CHAR(63) NOT NULL,
		Seen TIMESTAMP NOT NULL,
		Address VARCHAR(2048) NOT NULL
	)`)
	if err != nil {
		return err
	}

	row = db.QueryRow(`SELECT 'AddressesDeviceIDSeenIndex'::regclass`)
	if err = row.Scan(&tmp); err != nil {
		_, err = db.Exec(`CREATE INDEX AddressesDeviceIDSeenIndex ON Addresses (DeviceID, Seen)`)
	}
	if err != nil {
		return err
	}

	row = db.QueryRow(`SELECT 'AddressesDeviceIDAddressIndex'::regclass`)
	if err = row.Scan(&tmp); err != nil {
		_, err = db.Exec(`CREATE INDEX AddressesDeviceIDAddressIndex ON Addresses (DeviceID, Address)`)
	}
	if err != nil {
		return err
	}

	return nil
}
Exemplo n.º 3
0
// Init must be called once before DB is used.
//
// Since it used glog, make sure it is also called after flag.Parse is called.
func (c *DatabaseConfig) NewVersionedDB() (*VersionedDB, error) {
	if err := c.validate(); err != nil {
		return nil, err
	}

	// If there is a connection string then connect to the MySQL server.
	// This is for testing only. In production we get the relevant information
	// from the metadata server.
	var err error
	var DB *sql.DB = nil

	glog.Infoln("Opening SQL database.")
	DB, err = sql.Open(DEFAULT_DRIVER, c.MySQLString())
	if err != nil {
		return nil, fmt.Errorf("Failed to open connection to SQL server: %v", err)
	}

	glog.Infoln("Sending Ping.")
	if err := DB.Ping(); err != nil {
		return nil, fmt.Errorf("Failed to ping SQL server: %v", err)
	}

	// As outlined in this comment:
	// https://github.com/go-sql-driver/mysql/issues/257#issuecomment-48985975
	// We can remove this once we have determined it's not necessary.
	DB.SetMaxIdleConns(0)
	DB.SetMaxOpenConns(200)

	result := &VersionedDB{
		DB:             DB,
		migrationSteps: c.MigrationSteps,
	}

	// Make sure the migration table exists.
	if err := result.checkVersionTable(); err != nil {
		return nil, fmt.Errorf("Attempt to create version table returned: %v", err)
	}
	glog.Infoln("Version table OK.")

	// Ping the database occasionally to keep the connection fresh.
	go func() {
		c := time.Tick(1 * time.Minute)
		for _ = range c {
			if err := result.DB.Ping(); err != nil {
				glog.Warningln("Database failed to respond:", err)
			}
			glog.Infof("db: Successful ping")
		}
	}()

	return result, nil
}
Exemplo n.º 4
0
Arquivo: conn.go Projeto: Tecsisa/dex
func NewConnection(cfg Config) (*gorp.DbMap, error) {
	u, err := url.Parse(cfg.DSN)
	if err != nil {
		return nil, fmt.Errorf("parse DSN: %v", err)
	}
	var (
		db      *sql.DB
		dialect gorp.Dialect
	)
	switch u.Scheme {
	case "postgres":
		db, err = sql.Open("postgres", cfg.DSN)
		if err != nil {
			return nil, err
		}
		db.SetMaxIdleConns(cfg.MaxIdleConnections)
		db.SetMaxOpenConns(cfg.MaxOpenConnections)
		dialect = gorp.PostgresDialect{}
	case "sqlite3":
		db, err = sql.Open("sqlite3", u.Host)
		if err != nil {
			return nil, err
		}
		if u.Host == ":memory:" {
			// NOTE(ericchiang): sqlite3 coordinates concurrent clients through file locks.
			// In memory databases do not support concurrent calls. Limit the number of
			// open connections to 1.
			//
			// See: https://www.sqlite.org/faq.html#q5
			db.SetMaxOpenConns(1)
		}
		dialect = gorp.SqliteDialect{}
	default:
		return nil, errors.New("unrecognized database driver")
	}

	dbm := gorp.DbMap{Db: db, Dialect: dialect}

	for _, t := range tables {
		tm := dbm.AddTableWithName(t.model, t.name).SetKeys(t.autoinc, t.pkey...)
		for _, unique := range t.unique {
			cm := tm.ColMap(unique)
			if cm == nil {
				return nil, fmt.Errorf("no such column: %q", unique)
			}
			cm.SetUnique(true)
		}
	}
	return &dbm, nil
}
Exemplo n.º 5
0
func InitDB(cfg *Config) (err error) {
	if cfg.DBType == "sqlite" {
		cfg.DBType = "sqlite3"
	}
	dsn := cfg.DSN()
	if dsn == "" {
		return fmt.Errorf("unsupported database: `%s`", cfg.DBType)
	}

	var raw_db *sql.DB
	if raw_db, err = sql.Open(cfg.DBType, dsn); err != nil {
		return
	}
	if cfg.MaxIdleConns > 0 {
		raw_db.SetMaxIdleConns(cfg.MaxIdleConns)
	}
	if cfg.MaxOpenConns > 0 {
		raw_db.SetMaxOpenConns(cfg.MaxOpenConns)
	}

	var dialect gorp.Dialect
	switch cfg.DBType {
	case "mysql":
		dialect = gorp.MySQLDialect{Engine: "InnoDB", Encoding: "UTF-8"}
	case "postgres":
		dialect = gorp.PostgresDialect{}
	case "sqlite3":
		dialect = gorp.SqliteDialect{}
	}

	db = &gorp.DbMap{Db: raw_db, Dialect: dialect}
	if cfg.SqlTrace {
		db.TraceOn("[db]", logger{})
	}
	_tableLk.Lock()
	_tableFreezed = true
	_tableLk.Unlock()

	for _, table := range _tables {
		name, primary, model := table.Meta()
		t := db.AddTableWithName(model, name)
		if primary != "" {
			t.SetKeys(true, primary)
		}
	}
	return nil
}
Exemplo n.º 6
0
func InitMysqlConnPool() error {
	var err error
	var conn *sql.DB
	conn, err = sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=utf8&parseTime=true&loc=Local",
		GlobalConfig.MYSQL_USER, GlobalConfig.MYSQL_PASSWORD, GlobalConfig.MYSQL_ADDR, GlobalConfig.MYSQL_DBNAME))
	if err != nil {
		return err
	}
	err = conn.Ping()
	if err != nil {
		return err
	}
	conn.SetMaxIdleConns(GlobalConfig.MYSQL_MAX_IDLE_CONN)
	conn.SetMaxOpenConns(GlobalConfig.MYSQL_MAX_CONN)
	mydb = &db{conn}
	return nil
}
Exemplo n.º 7
0
func (dbSeq *SequenceDB) Open() (err error) {
	var db *sql.DB
	db, err = sql.Open("mysql", conf.Conf.SequenceDB.DSN)
	if err != nil {
		log.Printf("sequence db open error. %v", err)
		return err
	}

	err = db.Ping()
	if err != nil {
		log.Printf("sequence db ping error. %v", err)
		return err
	}

	db.SetMaxIdleConns(conf.Conf.SequenceDB.MaxIdleConns)
	db.SetMaxOpenConns(conf.Conf.SequenceDB.MaxOpenConns)

	dbSeq.db = db
	return nil
}
Exemplo n.º 8
0
func Db() (*sql.DB, error) {
	var (
		db  *sql.DB
		err error
	)

	config, err := BuildFromFile("../config/database.yml")
	if err != nil {
		log.Fatal(err)
	}

	db, err = sql.Open("postgres", "user="******" password="******" dbname="+config.DbName+" sslmode=disable")
	if err != nil {
		return nil, err
	}

	db.SetMaxIdleConns(10)
	db.SetMaxOpenConns(20)

	return db, nil
}
Exemplo n.º 9
0
Arquivo: db.go Projeto: RoyDong/potato
func NewDB() *sql.DB {
	var db *sql.DB
	var e error

	dsn := fmt.Sprintf("%s:%s@(%s:%d)/%s",
		C.User, C.Pass, C.Host, C.Port, C.DBname)

	if db, e = sql.Open(C.Type, dsn); e != nil {
		log.Fatal(e)
	}

	if C.MaxConn > 0 {
		db.SetMaxOpenConns(C.MaxConn)
	}

	if e = db.Ping(); e != nil {
		log.Fatal(e)
	}

	return db
}
Exemplo n.º 10
0
func (this *Conns) initSql(host, port, dbname, user, pass string, maxOpenConns, maxIdleConns int) (*sql.DB, error) {
	slimSqlLog("Init", "user:"******" pass:"******" host:"+host+" port:"+port+" dbname:"+dbname)
	var sqlDB *sql.DB
	var err error
	sqlDB, err = sql.Open("mysql", user+":"+pass+"@tcp("+host+":"+port+")/"+dbname+"?charset=utf8")
	if err != nil {
		return nil, err
	}
	if maxOpenConns > 0 {
		sqlDB.SetMaxOpenConns(maxOpenConns)
	} else {
		sqlDB.SetMaxOpenConns(2000)
	}
	if maxIdleConns > 0 {
		sqlDB.SetMaxIdleConns(maxIdleConns)
	} else {
		sqlDB.SetMaxIdleConns(1000)
	}
	err = sqlDB.Ping()
	if err != nil {
		return nil, err
	}
	return sqlDB, nil
}
Exemplo n.º 11
0
func main() {
	flag.Parse()
	var db *sql.DB
	var err error
	var url string
	if *usePostgres {
		if db, err = sql.Open("postgres", "dbname=postgres sslmode=disable"); err != nil {
			log.Fatal(err)
		}
	} else {
		security.SetReadFileFn(securitytest.Asset)
		serv := server.StartTestServer(nil)
		defer serv.Stop()
		url = "https://root@" + serv.ServingAddr() + "?certs=test_certs"
		if db, err = sql.Open("cockroach", url); err != nil {
			log.Fatal(err)
		}
	}
	if _, err := db.Exec("CREATE DATABASE bank"); err != nil {
		if pqErr, ok := err.(*pq.Error); *usePostgres && (!ok || pqErr.Code.Name() != "duplicate_database") {
			log.Fatal(err)
		}
	}
	db.Close()

	// Open db client with database settings.
	if *usePostgres {
		if db, err = sql.Open("postgres", "dbname=bank sslmode=disable"); err != nil {
			log.Fatal(err)
		}
	} else {
		if db, err = sql.Open("cockroach", url+"&database=bank"); err != nil {
			log.Fatal(err)
		}
	}

	// concurrency + 1, for this thread and the "concurrency" number of
	// goroutines that move money
	db.SetMaxOpenConns(*concurrency + 1)

	if _, err = db.Exec("CREATE TABLE IF NOT EXISTS accounts (id BIGINT PRIMARY KEY, balance BIGINT NOT NULL)"); err != nil {
		log.Fatal(err)
	}

	if _, err = db.Exec("TRUNCATE TABLE accounts"); err != nil {
		log.Fatal(err)
	}

	for i := 0; i < *numAccounts; i++ {
		if _, err = db.Exec("INSERT INTO accounts (id, balance) VALUES ($1, $2)", i, 1000); err != nil {
			log.Fatal(err)
		}
	}

	verifyBank(db)

	lastNow := time.Now()
	readings := make(chan measurement, 10000)

	for i := 0; i < *concurrency; i++ {
		go moveMoney(db, readings)
	}

	for range time.NewTicker(*balanceCheckInterval).C {
		now := time.Now()
		elapsed := time.Since(lastNow)
		lastNow = now
		transfers := len(readings)
		log.Infof("%d transfers were executed at %.1f/second.", transfers, float64(transfers)/elapsed.Seconds())
		if transfers > 0 {
			var aggr measurement
			for i := 0; i < transfers; i++ {
				reading := <-readings
				aggr.read += reading.read
				aggr.write += reading.write
				aggr.total += reading.total
			}
			d := time.Duration(transfers)
			log.Infof("read time: %v, write time: %v, txn time: %v", aggr.read/d, aggr.write/d, aggr.total/d)
		}
		verifyBank(db)
	}
}
Exemplo n.º 12
0
func main() {
	debugserver.AddFlags(flag.CommandLine)
	lagerflags.AddFlags(flag.CommandLine)
	etcdFlags := AddETCDFlags(flag.CommandLine)
	encryptionFlags := encryption.AddEncryptionFlags(flag.CommandLine)

	flag.Parse()

	cfhttp.Initialize(*communicationTimeout)

	logger, reconfigurableSink := lagerflags.New("bbs")
	logger.Info("starting")

	initializeDropsonde(logger)

	clock := clock.NewClock()

	consulClient, err := consuladapter.NewClientFromUrl(*consulCluster)
	if err != nil {
		logger.Fatal("new-consul-client-failed", err)
	}

	serviceClient := bbs.NewServiceClient(consulClient, clock)

	maintainer := initializeLockMaintainer(logger, serviceClient)

	_, portString, err := net.SplitHostPort(*listenAddress)
	if err != nil {
		logger.Fatal("failed-invalid-listen-address", err)
	}
	portNum, err := net.LookupPort("tcp", portString)
	if err != nil {
		logger.Fatal("failed-invalid-listen-port", err)
	}

	_, portString, err = net.SplitHostPort(*healthAddress)
	if err != nil {
		logger.Fatal("failed-invalid-health-address", err)
	}
	_, err = net.LookupPort("tcp", portString)
	if err != nil {
		logger.Fatal("failed-invalid-health-port", err)
	}

	registrationRunner := initializeRegistrationRunner(logger, consulClient, portNum, clock)

	var activeDB db.DB
	var sqlDB *sqldb.SQLDB
	var sqlConn *sql.DB
	var storeClient etcddb.StoreClient
	var etcdDB *etcddb.ETCDDB

	key, keys, err := encryptionFlags.Parse()
	if err != nil {
		logger.Fatal("cannot-setup-encryption", err)
	}
	keyManager, err := encryption.NewKeyManager(key, keys)
	if err != nil {
		logger.Fatal("cannot-setup-encryption", err)
	}
	cryptor := encryption.NewCryptor(keyManager, rand.Reader)

	etcdOptions, err := etcdFlags.Validate()
	if err != nil {
		logger.Fatal("etcd-validation-failed", err)
	}

	if etcdOptions.IsConfigured {
		storeClient = initializeEtcdStoreClient(logger, etcdOptions)
		etcdDB = initializeEtcdDB(logger, cryptor, storeClient, serviceClient, *desiredLRPCreationTimeout)
		activeDB = etcdDB
	}

	// If SQL database info is passed in, use SQL instead of ETCD
	if *databaseDriver != "" && *databaseConnectionString != "" {
		var err error
		connectionString := appendSSLConnectionStringParam(logger, *databaseDriver, *databaseConnectionString, *sqlCACertFile)

		sqlConn, err = sql.Open(*databaseDriver, connectionString)
		if err != nil {
			logger.Fatal("failed-to-open-sql", err)
		}
		defer sqlConn.Close()
		sqlConn.SetMaxOpenConns(*maxDatabaseConnections)
		sqlConn.SetMaxIdleConns(*maxDatabaseConnections)

		err = sqlConn.Ping()
		if err != nil {
			logger.Fatal("sql-failed-to-connect", err)
		}

		sqlDB = sqldb.NewSQLDB(sqlConn, *convergenceWorkers, *updateWorkers, format.ENCRYPTED_PROTO, cryptor, guidprovider.DefaultGuidProvider, clock, *databaseDriver)
		err = sqlDB.SetIsolationLevel(logger, sqldb.IsolationLevelReadCommitted)
		if err != nil {
			logger.Fatal("sql-failed-to-set-isolation-level", err)
		}

		err = sqlDB.CreateConfigurationsTable(logger)
		if err != nil {
			logger.Fatal("sql-failed-create-configurations-table", err)
		}
		activeDB = sqlDB
	}

	if activeDB == nil {
		logger.Fatal("no-database-configured", errors.New("no database configured"))
	}

	encryptor := encryptor.New(logger, activeDB, keyManager, cryptor, clock)

	migrationsDone := make(chan struct{})

	migrationManager := migration.NewManager(
		logger,
		etcdDB,
		storeClient,
		sqlDB,
		sqlConn,
		cryptor,
		migrations.Migrations,
		migrationsDone,
		clock,
		*databaseDriver,
	)

	desiredHub := events.NewHub()
	actualHub := events.NewHub()

	repTLSConfig := &rep.TLSConfig{
		RequireTLS:      *repRequireTLS,
		CaCertFile:      *repCACert,
		CertFile:        *repClientCert,
		KeyFile:         *repClientKey,
		ClientCacheSize: *repClientSessionCacheSize,
	}

	httpClient := cfhttp.NewClient()
	repClientFactory, err := rep.NewClientFactory(httpClient, httpClient, repTLSConfig)
	if err != nil {
		logger.Fatal("new-rep-client-factory-failed", err)
	}

	auctioneerClient := initializeAuctioneerClient(logger)

	exitChan := make(chan struct{})

	var accessLogger lager.Logger
	if *accessLogPath != "" {
		accessLogger = lager.NewLogger("bbs-access")
		file, err := os.OpenFile(*accessLogPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			logger.Error("invalid-access-log-path", err, lager.Data{"access-log-path": *accessLogPath})
			os.Exit(1)
		}
		accessLogger.RegisterSink(lager.NewWriterSink(file, lager.INFO))
	}

	var tlsConfig *tls.Config
	if *requireSSL {
		tlsConfig, err = cfhttp.NewTLSConfig(*certFile, *keyFile, *caFile)
		if err != nil {
			logger.Fatal("tls-configuration-failed", err)
		}
	}

	cbWorkPool := taskworkpool.New(logger, *taskCallBackWorkers, taskworkpool.HandleCompletedTask, tlsConfig)

	handler := handlers.New(
		logger,
		accessLogger,
		*updateWorkers,
		*convergenceWorkers,
		activeDB,
		desiredHub,
		actualHub,
		cbWorkPool,
		serviceClient,
		auctioneerClient,
		repClientFactory,
		migrationsDone,
		exitChan,
	)

	metricsNotifier := metrics.NewPeriodicMetronNotifier(logger)

	retirer := controllers.NewActualLRPRetirer(activeDB, actualHub, repClientFactory, serviceClient)
	lrpConvergenceController := controllers.NewLRPConvergenceController(logger, activeDB, actualHub, auctioneerClient, serviceClient, retirer, *convergenceWorkers)
	taskController := controllers.NewTaskController(activeDB, cbWorkPool, auctioneerClient, serviceClient, repClientFactory)

	convergerProcess := converger.New(
		logger,
		clock,
		lrpConvergenceController,
		taskController,
		serviceClient,
		*convergeRepeatInterval,
		*kickTaskDuration,
		*expirePendingTaskDuration,
		*expireCompletedTaskDuration)

	var server ifrit.Runner
	if tlsConfig != nil {
		server = http_server.NewTLSServer(*listenAddress, handler, tlsConfig)
	} else {
		server = http_server.New(*listenAddress, handler)
	}

	healthcheckServer := http_server.New(*healthAddress, http.HandlerFunc(healthCheckHandler))

	members := grouper.Members{
		{"healthcheck", healthcheckServer},
		{"lock-maintainer", maintainer},
		{"workpool", cbWorkPool},
		{"server", server},
		{"migration-manager", migrationManager},
		{"encryptor", encryptor},
		{"hub-maintainer", hubMaintainer(logger, desiredHub, actualHub)},
		{"metrics", *metricsNotifier},
		{"converger", convergerProcess},
		{"registration-runner", registrationRunner},
	}

	if dbgAddr := debugserver.DebugAddress(flag.CommandLine); dbgAddr != "" {
		members = append(grouper.Members{
			{"debug-server", debugserver.Runner(dbgAddr, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group))
	go func() {
		// If a handler writes to this channel, we've hit an unrecoverable error
		// and should shut down (cleanly)
		<-exitChan
		monitor.Signal(os.Interrupt)
	}()

	logger.Info("started")

	err = <-monitor.Wait()
	if sqlConn != nil {
		sqlConn.Close()
	}
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}
Exemplo n.º 13
0
func main() {
	flag.Parse()
	var db *sql.DB
	var err error
	var url string
	if *usePostgres {
		if db, err = sql.Open("postgres", "dbname=postgres sslmode=disable"); err != nil {
			log.Fatal(err)
		}
	} else {
		security.SetReadFileFn(securitytest.Asset)
		serv := server.StartTestServer(nil)
		defer serv.Stop()
		url = "https://root@" + serv.ServingAddr() + "?certs=test_certs"
		if db, err = sql.Open("cockroach", url); err != nil {
			log.Fatal(err)
		}
	}
	if _, err := db.Exec("CREATE DATABASE bank"); err != nil {
		if pqErr, ok := err.(*pq.Error); *usePostgres && (!ok || pqErr.Code.Name() != "duplicate_database") {
			log.Fatal(err)
		}
	}
	db.Close()

	// Open db client with database settings.
	if *usePostgres {
		if db, err = sql.Open("postgres", "dbname=bank sslmode=disable"); err != nil {
			log.Fatal(err)
		}
	} else {
		if db, err = sql.Open("cockroach", url+"&database=bank"); err != nil {
			log.Fatal(err)
		}
	}

	// concurrency + 1, for this thread and the "concurrency" number of
	// goroutines that move money
	db.SetMaxOpenConns(*concurrency + 1)

	if _, err = db.Exec("CREATE TABLE IF NOT EXISTS accounts (id BIGINT PRIMARY KEY, balance BIGINT NOT NULL)"); err != nil {
		log.Fatal(err)
	}

	if _, err = db.Exec("TRUNCATE TABLE accounts"); err != nil {
		log.Fatal(err)
	}

	for i := 0; i < *numAccounts; i++ {
		if _, err = db.Exec("INSERT INTO accounts (id, balance) VALUES ($1, $2)", i, 0); err != nil {
			log.Fatal(err)
		}
	}

	verifyBank(db)

	var lastNumTransfers uint64
	lastNow := time.Now()

	for i := 0; i < *concurrency; i++ {
		go moveMoney(db)
	}

	for range time.NewTicker(*balanceCheckInterval).C {
		now := time.Now()
		elapsed := time.Since(lastNow)
		numTransfers := atomic.LoadUint64(&numTransfers)
		log.Infof("%d transfers were executed at %.1f/second.", (numTransfers - lastNumTransfers), float64(numTransfers-lastNumTransfers)/elapsed.Seconds())
		verifyBank(db)
		lastNumTransfers = numTransfers
		lastNow = now
	}
}
Exemplo n.º 14
0
func main() {
	flag.Usage = usage
	flag.Parse()

	if flag.NArg() != 1 {
		usage()
		os.Exit(2)
	}

	dbURL := flag.Arg(0)

	isPostgres := false
	var db *sql.DB
	var err error
	var parsedURL *url.URL

	if *dbDriver == "postgres" {
		isPostgres = true
	} else if *dbDriver != "cockroach" {
		log.Fatal("--driver only supports \"cockroach\" and \"postgres\"")
	}

	if parsedURL, err = url.Parse(dbURL); err != nil {
		log.Fatal(err)
	}

	if isPostgres {
		if db, err = sql.Open("postgres", parsedURL.String()); err != nil {
			log.Fatal(err)
		}
	} else {
		if db, err = sql.Open("cockroach", parsedURL.String()); err != nil {
			log.Fatal(err)
		}
	}

	if _, err := db.Exec("CREATE DATABASE bank"); err != nil {
		if pqErr, ok := err.(*pq.Error); isPostgres && (!ok || pqErr.Code.Name() != "duplicate_database") {
			log.Fatal(err)
		}
	}
	if err := db.Close(); err != nil {
		log.Fatal(err)
	}

	// Open db client with database settings.
	if isPostgres {
		q := parsedURL.Query()
		q.Set("dbname", "bank")
		parsedURL.RawQuery = q.Encode()
		if db, err = sql.Open("postgres", parsedURL.String()); err != nil {
			log.Fatal(err)
		}
	} else {
		q := parsedURL.Query()
		q.Set("database", "bank")
		parsedURL.RawQuery = q.Encode()
		if db, err = sql.Open("cockroach", parsedURL.String()); err != nil {
			log.Fatal(err)
		}
	}

	// concurrency + 1, for this thread and the "concurrency" number of
	// goroutines that move money
	db.SetMaxOpenConns(*concurrency + 1)

	if _, err = db.Exec("CREATE TABLE IF NOT EXISTS accounts (id BIGINT PRIMARY KEY, balance BIGINT NOT NULL)"); err != nil {
		log.Fatal(err)
	}

	if _, err = db.Exec("TRUNCATE TABLE accounts"); err != nil {
		log.Fatal(err)
	}

	for i := 0; i < *numAccounts; i++ {
		if _, err = db.Exec("INSERT INTO accounts (id, balance) VALUES ($1, $2)", i, 1000); err != nil {
			log.Fatal(err)
		}
	}

	verifyBank(db)

	lastNow := time.Now()
	readings := make(chan measurement, 10000)

	for i := 0; i < *concurrency; i++ {
		go moveMoney(db, readings)
	}

	for range time.NewTicker(*balanceCheckInterval).C {
		now := time.Now()
		elapsed := time.Since(lastNow)
		lastNow = now
		transfers := len(readings)
		log.Printf("%d transfers were executed at %.1f/second.", transfers, float64(transfers)/elapsed.Seconds())
		if transfers > 0 {
			var aggr measurement
			for i := 0; i < transfers; i++ {
				reading := <-readings
				aggr.read += reading.read
				aggr.write += reading.write
				aggr.total += reading.total
			}
			d := time.Duration(transfers)
			log.Printf("read time: %v, write time: %v, txn time: %v", aggr.read/d, aggr.write/d, aggr.total/d)
		}
		verifyBank(db)
	}
}
Exemplo n.º 15
0
func (a *Attacker) hit(tr Targeter, tm time.Time) *Result {
	var (
		res = Result{Timestamp: tm}
		tgt Target
		err error
	)

	defer func() {
		res.Latency = time.Since(tm)
		if err != nil {
			res.Error = err.Error()
		}
	}()

	if err = tr(&tgt); err != nil {
		return &res
	}

	req, err := tgt.Query()
	if err != nil {
		return &res
	}

	/*
		a.cnn, err = sql.Open("mysql", a.dsn)
		if err != nil {
			res.Code = 500
			res.Error = err.Error()
			return &res
		}
		defer a.cnn.Close()
		a.cnn.SetMaxIdleConns(a.maxIdleConns)
		a.cnn.SetMaxOpenConns(a.maxOpenConns)
	*/
	//log.Printf("query:%s", req)

	var r *sql.Rows
	if a.persistent {
		r, err = a.cnn.Query(req)
	} else {
		var cnn *sql.DB
		cnn, err = sql.Open("mysql", a.dsn)
		cnn.SetMaxIdleConns(0)
		cnn.SetMaxOpenConns(1)
		cnn.SetConnMaxLifetime(1 * time.Nanosecond)
		if err != nil {
			res.Code = 500
			res.Error = err.Error()
			return &res
		}
		r, err = cnn.Query(req)
	}
	if err != nil {
		// ignore redirect errors when the user set --redirects=NoFollow
		if a.redirects == NoFollow && strings.Contains(err.Error(), "stopped after") {
			err = nil
		}
		return &res
	}
	defer r.Close()
	num := 0
	for r.Next() {
		num++
		/*
			var id interface{}
				if err := r.Scan(&id); err != nil {
					res.Code = 501
					res.Error = err.Error()
					res.Error = fmt.Sprintf("row scan err:%s: query:%s", err.Error(), req)
					return &res
				}
		*/
	}
	//fmt.Fprintf(os.Stderr, "%d,", num)
	res.BytesIn = 0
	res.BytesOut = 0

	err = r.Err()

	if err == nil {
		if num == 0 {
			res.Code = 201
		} else {
			res.Code = 200
		}
	} else {
		res.Code = 500
		res.Error = fmt.Sprintf("%s: query:%s", err.Error(), req)
	}
	return &res
}