guerrilla_db_redis.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470
  1. package backends
  2. // This backend is presented here as an example only, please modify it to your needs.
  3. //
  4. // Deprecated: as of 14th Feb 2017, backends are composed via config, by chaining Processors (files prefixed with p_*)
  5. //
  6. // The backend stores the email data in Redis.
  7. // Other meta-information is stored in MySQL to be joined later.
  8. // A lot of email gets discarded without viewing on Guerrilla Mail,
  9. // so it's much faster to put in Redis, where other programs can
  10. // process it later, without touching the disk.
  11. //
  12. // Some features:
  13. // - It batches the SQL inserts into a single query and inserts either after a time threshold or if the batch is full
  14. // - If the mysql driver crashes, it's able to recover, log the incident and resume again.
  15. // - It also does a clean shutdown - it tries to save everything before returning
  16. //
  17. // Short history:
  18. // Started with issuing an insert query for each single email and another query to update the tally
  19. // Then applied the following optimizations:
  20. // - Moved tally updates to another background process which does the tallying in a single query
  21. // - Changed the MySQL queries to insert in batch
  22. // - Made a Compressor that recycles buffers using sync.Pool
  23. // The result was around 400% speed improvement. If you know of any more improvements, please share!
  24. // - Added the recovery mechanism,
  25. import (
  26. "fmt"
  27. "time"
  28. "github.com/garyburd/redigo/redis"
  29. "bytes"
  30. "compress/zlib"
  31. "database/sql"
  32. "github.com/flashmob/go-guerrilla/envelope"
  33. "github.com/go-sql-driver/mysql"
  34. "io"
  35. "runtime/debug"
  36. "strings"
  37. "sync"
  38. )
  39. // how many rows to batch at a time
  40. const GuerrillaDBAndRedisBatchMax = 2
  41. // tick on every...
  42. const GuerrillaDBAndRedisBatchTimeout = time.Second * 3
  43. func init() {
  44. backends["guerrilla-db-redis"] = &GuerrillaDBAndRedisBackend{}
  45. }
  46. type GuerrillaDBAndRedisBackend struct {
  47. config guerrillaDBAndRedisConfig
  48. batcherWg sync.WaitGroup
  49. // cache prepared queries
  50. cache stmtCache
  51. }
  52. // statement cache. It's an array, not slice
  53. type stmtCache [GuerrillaDBAndRedisBatchMax]*sql.Stmt
  54. type guerrillaDBAndRedisConfig struct {
  55. NumberOfWorkers int `json:"save_workers_size"`
  56. MysqlTable string `json:"mail_table"`
  57. MysqlDB string `json:"mysql_db"`
  58. MysqlHost string `json:"mysql_host"`
  59. MysqlPass string `json:"mysql_pass"`
  60. MysqlUser string `json:"mysql_user"`
  61. RedisExpireSeconds int `json:"redis_expire_seconds"`
  62. RedisInterface string `json:"redis_interface"`
  63. PrimaryHost string `json:"primary_mail_host"`
  64. }
  65. func convertError(name string) error {
  66. return fmt.Errorf("failed to load backend config (%s)", name)
  67. }
  68. // Load the backend config for the backend. It has already been unmarshalled
  69. // from the main config file 'backend' config "backend_config"
  70. // Now we need to convert each type and copy into the guerrillaDBAndRedisConfig struct
  71. func (g *GuerrillaDBAndRedisBackend) loadConfig(backendConfig BackendConfig) (err error) {
  72. configType := baseConfig(&guerrillaDBAndRedisConfig{})
  73. bcfg, err := Service.extractConfig(backendConfig, configType)
  74. if err != nil {
  75. return err
  76. }
  77. m := bcfg.(*guerrillaDBAndRedisConfig)
  78. g.config = *m
  79. return nil
  80. }
  81. func (g *GuerrillaDBAndRedisBackend) getNumberOfWorkers() int {
  82. return g.config.NumberOfWorkers
  83. }
  84. type redisClient struct {
  85. isConnected bool
  86. conn redis.Conn
  87. time int
  88. }
  89. // compressedData struct will be compressed using zlib when printed via fmt
  90. type compressedData struct {
  91. extraHeaders []byte
  92. data *bytes.Buffer
  93. pool *sync.Pool
  94. }
  95. // newCompressedData returns a new CompressedData
  96. func newCompressedData() *compressedData {
  97. var p = sync.Pool{
  98. New: func() interface{} {
  99. var b bytes.Buffer
  100. return &b
  101. },
  102. }
  103. return &compressedData{
  104. pool: &p,
  105. }
  106. }
  107. // Set the extraheaders and buffer of data to compress
  108. func (c *compressedData) set(b []byte, d *bytes.Buffer) {
  109. c.extraHeaders = b
  110. c.data = d
  111. }
  112. // implement Stringer interface
  113. func (c *compressedData) String() string {
  114. if c.data == nil {
  115. return ""
  116. }
  117. //borrow a buffer form the pool
  118. b := c.pool.Get().(*bytes.Buffer)
  119. // put back in the pool
  120. defer func() {
  121. b.Reset()
  122. c.pool.Put(b)
  123. }()
  124. var r *bytes.Reader
  125. w, _ := zlib.NewWriterLevel(b, zlib.BestSpeed)
  126. r = bytes.NewReader(c.extraHeaders)
  127. io.Copy(w, r)
  128. io.Copy(w, c.data)
  129. w.Close()
  130. return b.String()
  131. }
  132. // clear it, without clearing the pool
  133. func (c *compressedData) clear() {
  134. c.extraHeaders = []byte{}
  135. c.data = nil
  136. }
  137. // prepares the sql query with the number of rows that can be batched with it
  138. func (g *GuerrillaDBAndRedisBackend) prepareInsertQuery(rows int, db *sql.DB) *sql.Stmt {
  139. if rows == 0 {
  140. panic("rows argument cannot be 0")
  141. }
  142. if g.cache[rows-1] != nil {
  143. return g.cache[rows-1]
  144. }
  145. sqlstr := "INSERT INTO " + g.config.MysqlTable + " "
  146. sqlstr += "(`date`, `to`, `from`, `subject`, `body`, `charset`, `mail`, `spam_score`, `hash`, `content_type`, `recipient`, `has_attach`, `ip_addr`, `return_path`, `is_tls`)"
  147. sqlstr += " values "
  148. values := "(NOW(), ?, ?, ?, ? , 'UTF-8' , ?, 0, ?, '', ?, 0, ?, ?, ?)"
  149. // add more rows
  150. comma := ""
  151. for i := 0; i < rows; i++ {
  152. sqlstr += comma + values
  153. if comma == "" {
  154. comma = ","
  155. }
  156. }
  157. stmt, sqlErr := db.Prepare(sqlstr)
  158. if sqlErr != nil {
  159. Log().WithError(sqlErr).Fatalf("failed while db.Prepare(INSERT...)")
  160. }
  161. // cache it
  162. g.cache[rows-1] = stmt
  163. return stmt
  164. }
  165. func (g *GuerrillaDBAndRedisBackend) doQuery(c int, db *sql.DB, insertStmt *sql.Stmt, vals *[]interface{}) {
  166. var execErr error
  167. defer func() {
  168. if r := recover(); r != nil {
  169. //logln(1, fmt.Sprintf("Recovered in %v", r))
  170. Log().Error("Recovered form panic:", r, string(debug.Stack()))
  171. sum := 0
  172. for _, v := range *vals {
  173. if str, ok := v.(string); ok {
  174. sum = sum + len(str)
  175. }
  176. }
  177. Log().Errorf("panic while inserting query [%s] size:%d, err %v", r, sum, execErr)
  178. panic("query failed")
  179. }
  180. }()
  181. // prepare the query used to insert when rows reaches batchMax
  182. insertStmt = g.prepareInsertQuery(c, db)
  183. _, execErr = insertStmt.Exec(*vals...)
  184. if execErr != nil {
  185. Log().WithError(execErr).Error("There was a problem the insert")
  186. }
  187. }
  188. // Batches the rows from the feeder chan in to a single INSERT statement.
  189. // Execute the batches query when:
  190. // - number of batched rows reaches a threshold, i.e. count n = threshold
  191. // - or, no new rows within a certain time, i.e. times out
  192. // The goroutine can either exit if there's a panic or feeder channel closes
  193. // it returns feederOk which signals if the feeder chanel was ok (still open) while returning
  194. // if it feederOk is false, then it means the feeder chanel is closed
  195. func (g *GuerrillaDBAndRedisBackend) insertQueryBatcher(feeder chan []interface{}, db *sql.DB) (feederOk bool) {
  196. // controls shutdown
  197. defer g.batcherWg.Done()
  198. g.batcherWg.Add(1)
  199. // vals is where values are batched to
  200. var vals []interface{}
  201. // how many rows were batched
  202. count := 0
  203. // The timer will tick every second.
  204. // Interrupting the select clause when there's no data on the feeder channel
  205. t := time.NewTimer(GuerrillaDBAndRedisBatchTimeout)
  206. // prepare the query used to insert when rows reaches batchMax
  207. insertStmt := g.prepareInsertQuery(GuerrillaDBAndRedisBatchMax, db)
  208. // inserts executes a batched insert query, clears the vals and resets the count
  209. insert := func(c int) {
  210. if c > 0 {
  211. g.doQuery(c, db, insertStmt, &vals)
  212. }
  213. vals = nil
  214. count = 0
  215. }
  216. defer func() {
  217. if r := recover(); r != nil {
  218. Log().Error("insertQueryBatcher caught a panic", r)
  219. }
  220. }()
  221. // Keep getting values from feeder and add to batch.
  222. // if feeder times out, execute the batched query
  223. // otherwise, execute the batched query once it reaches the GuerrillaDBAndRedisBatchMax threshold
  224. feederOk = true
  225. for {
  226. select {
  227. // it may panic when reading on a closed feeder channel. feederOK detects if it was closed
  228. case row, feederOk := <-feeder:
  229. if row == nil {
  230. Log().Info("Query batchaer exiting")
  231. // Insert any remaining rows
  232. insert(count)
  233. return feederOk
  234. }
  235. vals = append(vals, row...)
  236. count++
  237. Log().Debug("new feeder row:", row, " cols:", len(row), " count:", count, " worker", workerId)
  238. if count >= GuerrillaDBAndRedisBatchMax {
  239. insert(GuerrillaDBAndRedisBatchMax)
  240. }
  241. // stop timer from firing (reset the interrupt)
  242. if !t.Stop() {
  243. <-t.C
  244. }
  245. t.Reset(GuerrillaDBAndRedisBatchTimeout)
  246. case <-t.C:
  247. // anything to insert?
  248. if n := len(vals); n > 0 {
  249. insert(count)
  250. }
  251. t.Reset(GuerrillaDBAndRedisBatchTimeout)
  252. }
  253. }
  254. }
  255. func trimToLimit(str string, limit int) string {
  256. ret := strings.TrimSpace(str)
  257. if len(str) > limit {
  258. ret = str[:limit]
  259. }
  260. return ret
  261. }
  262. var workerId = 0
  263. func (g *GuerrillaDBAndRedisBackend) mysqlConnect() (*sql.DB, error) {
  264. conf := mysql.Config{
  265. User: g.config.MysqlUser,
  266. Passwd: g.config.MysqlPass,
  267. DBName: g.config.MysqlDB,
  268. Net: "tcp",
  269. Addr: g.config.MysqlHost,
  270. ReadTimeout: GuerrillaDBAndRedisBatchTimeout + (time.Second * 10),
  271. WriteTimeout: GuerrillaDBAndRedisBatchTimeout + (time.Second * 10),
  272. Params: map[string]string{"collation": "utf8_general_ci"},
  273. }
  274. if db, err := sql.Open("mysql", conf.FormatDSN()); err != nil {
  275. Log().Error("cannot open mysql", err)
  276. return nil, err
  277. } else {
  278. return db, nil
  279. }
  280. }
  281. func (g *GuerrillaDBAndRedisBackend) saveMailWorker(saveMailChan chan *savePayload) {
  282. var to, body string
  283. var redisErr error
  284. workerId++
  285. redisClient := &redisClient{}
  286. var db *sql.DB
  287. var err error
  288. db, err = g.mysqlConnect()
  289. if err != nil {
  290. Log().Fatalf("cannot open mysql: %s", err)
  291. }
  292. // start the query SQL batching where we will send data via the feeder channel
  293. feeder := make(chan []interface{}, 1)
  294. go func() {
  295. for {
  296. if feederOK := g.insertQueryBatcher(feeder, db); !feederOK {
  297. Log().Debug("insertQueryBatcher exited")
  298. return
  299. }
  300. // if insertQueryBatcher panics, it can recover and go in again
  301. Log().Debug("resuming insertQueryBatcher")
  302. }
  303. }()
  304. defer func() {
  305. if r := recover(); r != nil {
  306. //recover form closed channel
  307. Log().Error("panic recovered in saveMailWorker", r)
  308. }
  309. db.Close()
  310. if redisClient.conn != nil {
  311. Log().Infof("closed redis")
  312. redisClient.conn.Close()
  313. }
  314. // close the feeder & wait for query batcher to exit.
  315. close(feeder)
  316. g.batcherWg.Wait()
  317. }()
  318. var vals []interface{}
  319. data := newCompressedData()
  320. // receives values from the channel repeatedly until it is closed.
  321. for {
  322. payload := <-saveMailChan
  323. if payload == nil {
  324. Log().Debug("No more saveMailChan payload")
  325. return
  326. }
  327. Log().Debug("Got mail from chan", payload.mail.RemoteAddress)
  328. to = trimToLimit(strings.TrimSpace(payload.mail.RcptTo[0].User)+"@"+g.config.PrimaryHost, 255)
  329. payload.mail.Helo = trimToLimit(payload.mail.Helo, 255)
  330. host := trimToLimit(payload.mail.RcptTo[0].Host, 255)
  331. ts := fmt.Sprintf("%d", time.Now().UnixNano())
  332. payload.mail.ParseHeaders()
  333. hash := MD5Hex(
  334. to,
  335. payload.mail.MailFrom.String(),
  336. payload.mail.Subject,
  337. ts)
  338. // Add extra headers
  339. var addHead string
  340. addHead += "Delivered-To: " + to + "\r\n"
  341. addHead += "Received: from " + payload.mail.Helo + " (" + payload.mail.Helo + " [" + payload.mail.RemoteAddress + "])\r\n"
  342. addHead += " by " + host + " with SMTP id " + hash + "@" + host + ";\r\n"
  343. addHead += " " + time.Now().Format(time.RFC1123Z) + "\r\n"
  344. // data will be compressed when printed, with addHead added to beginning
  345. data.set([]byte(addHead), &payload.mail.Data)
  346. body = "gzencode"
  347. // data will be written to redis - it implements the Stringer interface, redigo uses fmt to
  348. // print the data to redis.
  349. redisErr = redisClient.redisConnection(g.config.RedisInterface)
  350. if redisErr == nil {
  351. _, doErr := redisClient.conn.Do("SETEX", hash, g.config.RedisExpireSeconds, data)
  352. if doErr == nil {
  353. body = "redis" // the backend system will know to look in redis for the message data
  354. data.clear() // blank
  355. }
  356. } else {
  357. Log().WithError(redisErr).Warn("Error while connecting redis")
  358. }
  359. vals = []interface{}{} // clear the vals
  360. vals = append(vals,
  361. trimToLimit(to, 255),
  362. trimToLimit(payload.mail.MailFrom.String(), 255),
  363. trimToLimit(payload.mail.Subject, 255),
  364. body,
  365. data.String(),
  366. hash,
  367. trimToLimit(to, 255),
  368. payload.mail.RemoteAddress,
  369. trimToLimit(payload.mail.MailFrom.String(), 255),
  370. payload.mail.TLS)
  371. feeder <- vals
  372. payload.savedNotify <- &saveStatus{nil, hash}
  373. }
  374. }
  375. func (c *redisClient) redisConnection(redisInterface string) (err error) {
  376. if c.isConnected == false {
  377. c.conn, err = redis.Dial("tcp", redisInterface)
  378. if err != nil {
  379. // handle error
  380. return err
  381. }
  382. c.isConnected = true
  383. }
  384. return nil
  385. }
  386. // test database connection settings
  387. func (g *GuerrillaDBAndRedisBackend) testSettings() (err error) {
  388. var db *sql.DB
  389. if db, err = g.mysqlConnect(); err != nil {
  390. err = fmt.Errorf("MySql cannot connect, check your settings: %s", err)
  391. } else {
  392. db.Close()
  393. }
  394. redisClient := &redisClient{}
  395. if redisErr := redisClient.redisConnection(g.config.RedisInterface); redisErr != nil {
  396. err = fmt.Errorf("Redis cannot connect, check your settings: %s", redisErr)
  397. }
  398. return
  399. }
  400. func (g *GuerrillaDBAndRedisBackend) Initialize(config BackendConfig) error {
  401. err := g.loadConfig(config)
  402. if err != nil {
  403. return err
  404. }
  405. return nil
  406. }
  407. // does nothing
  408. func (g *GuerrillaDBAndRedisBackend) Process(mail *envelope.Envelope) BackendResult {
  409. return NewBackendResult("250 OK")
  410. }
  411. // does nothing
  412. func (g *GuerrillaDBAndRedisBackend) Shutdown() error {
  413. return nil
  414. }