func createGorm(setting *Setting) (DB, error) {
db, err := gorm.Open(
setting.Dialect(),
setting.URL,
)
if err != nil {
return nil, err
}
log.Info("Open Database !")
db.SingularTable(true)
if log.IsDebugEnabled() {
db.LogMode(true)
}
db.DB()
if setting.MaxIdleConns > 0 {
db.DB().SetMaxIdleConns(setting.MaxIdleConns)
}
if setting.MaxOpenConns > 0 {
db.DB().SetMaxOpenConns(setting.MaxOpenConns)
}
return &gormDB{
Dbm: &db,
}, nil
}
func (b *BridgeDB) Start(converter Converter) error {
if err := b.mqttClient.Connect(); err != nil {
log.Error(err)
return err
}
defer b.Stop()
var st time.Time
var counterLimit int64 = 100000
sem := make(chan int, b.setting.Concurrency)
for {
err := b.mqttClient.Subscribe(
b.setting.Topic,
mqtt.QOS_ZERO,
func(msg mqtt.Message) error {
atomic.AddInt64(&b.counter, 1)
if b.counter == 1 {
log.Infof(
"Benchmark mqtt-subscriber pg-bridge start...",
)
st = time.Now()
}
log.Infof("counter=%v", b.counter)
sem <- 1
b.wg.Add(1)
go func(topic string, msgData []byte) {
log.Info(string(msgData))
r := converter(topic)
json.Unmarshal(msgData, r)
err := b.dbm.Save(r)
if err != nil {
log.Error(err)
}
b.wg.Done()
<-sem
}(msg.Topic(), msg.Payload())
return nil
},
)
if err != nil {
log.Error(err)
break
}
if b.counter >= counterLimit {
break
}
time.Sleep(time.Millisecond)
}
b.wg.Wait()
log.Infof("counter=%v", b.counter)
log.Infof(
"Benchmark mqtt-subscriber pg-bridge end.[elapsedtime=%f sec]",
time.Now().Sub(st).Seconds(),
)
return nil
}
func main() {
initSchema()
setting := &bdb.Setting{
Topic: env.GetEnv("MQTT_TOPIC", "test/#").String(),
Concurrency: int(env.GetEnv("MQTT_CONCURRENCY", "10").Int64()),
DB: &db.Setting{
URL: os.Getenv("DATABASE_URL"),
MaxIdleConns: int(env.GetEnv("DATABASE_MAX_IDLE_CONNS", "10").Int64()),
MaxOpenConns: int(env.GetEnv("DATABASE_MAX_OPEN_CONNS", "10").Int64()),
},
MQTT: &mqtt.Setting{
ClientId: fmt.Sprintf("Bridge-DB-%v", time.Now().Unix()),
URL: env.GetEnv("MQTT_BROKER_URL", "tcp://localhost:1883").String(),
},
}
bridgeDB, err := bdb.Create(setting)
if err != nil {
panic(err)
}
go bridgeDB.Start(
func(topic string) interface{} {
return &RawData{}
},
)
api := rest.NewApi()
api.Use(rest.DefaultCommonStack...)
router, err := rest.MakeRouter(
&rest.Route{
HttpMethod: "GET",
PathExp: "/",
Func: func(w rest.ResponseWriter, r *rest.Request) {
w.WriteJson(OK)
},
},
)
if err != nil {
panic(err)
}
api.SetApp(router)
server := &graceful.Server{
Server: &http.Server{
Addr: ":" + env.GetEnv("PORT", "9000").String(),
Handler: api.MakeHandler(),
},
Timeout: 10 * time.Second,
NoSignalHandling: false,
ShutdownInitiated: func() {
log.Info("Call ShutdownInitiated")
bridgeDB.Stop()
},
}
server.ListenAndServe()
}
func (db *sqlxDB) CloseDB() error {
if db == nil || db.Dbm == nil {
return fmt.Errorf("Already closed db")
}
if err := db.Dbm.Close(); err != nil {
log.Warn("Failed to close Database !")
} else {
log.Info("Close Database !")
}
return nil
}
func createSqlx(setting *Setting) (DB, error) {
db, err := sqlx.Connect(
setting.Dialect(),
setting.URL,
)
if err != nil {
return nil, err
}
log.Info("Open Database !")
if setting.MaxIdleConns > 0 {
db.DB.SetMaxIdleConns(setting.MaxIdleConns)
}
if setting.MaxOpenConns > 0 {
db.DB.SetMaxOpenConns(setting.MaxOpenConns)
}
return &sqlxDB{
Dbm: db,
}, nil
}