p_guerrilla_db_redis.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493
  1. package backends
  2. import (
  3. "bytes"
  4. "compress/zlib"
  5. "database/sql"
  6. "fmt"
  7. "io"
  8. "math/rand"
  9. "runtime/debug"
  10. "strings"
  11. "sync"
  12. "time"
  13. "github.com/flashmob/go-guerrilla/mail"
  14. )
  15. // ----------------------------------------------------------------------------------
  16. // Processor Name: GuerrillaRedisDB
  17. // ----------------------------------------------------------------------------------
  18. // Description : Saves the body to redis, meta data to SQL. Example only.
  19. // : Limitation: it doesn't save multiple recipients or validate them
  20. // ----------------------------------------------------------------------------------
  21. // Config Options: ...
  22. // --------------:-------------------------------------------------------------------
  23. // Input : envelope
  24. // ----------------------------------------------------------------------------------
  25. // Output :
  26. // ----------------------------------------------------------------------------------
  27. func init() {
  28. processors["guerrillaredisdb"] = func() Decorator {
  29. return GuerrillaDbRedis()
  30. }
  31. }
  32. var queryBatcherId = 0
  33. // how many rows to batch at a time
  34. const GuerrillaDBAndRedisBatchMax = 50
  35. // tick on every...
  36. const GuerrillaDBAndRedisBatchTimeout = time.Second * 3
  37. type GuerrillaDBAndRedisBackend struct {
  38. config *guerrillaDBAndRedisConfig
  39. batcherWg sync.WaitGroup
  40. // cache prepared queries
  41. cache stmtCache
  42. batcherStoppers []chan bool
  43. }
  44. // statement cache. It's an array, not slice
  45. type stmtCache [GuerrillaDBAndRedisBatchMax]*sql.Stmt
  46. type guerrillaDBAndRedisConfig struct {
  47. NumberOfWorkers int `json:"save_workers_size"`
  48. Table string `json:"mail_table"`
  49. Driver string `json:"sql_driver"`
  50. DSN string `json:"sql_dsn"`
  51. RedisExpireSeconds int `json:"redis_expire_seconds"`
  52. RedisInterface string `json:"redis_interface"`
  53. PrimaryHost string `json:"primary_mail_host"`
  54. BatchTimeout int `json:"redis_sql_batch_timeout,omitempty"`
  55. }
  56. // Load the backend config for the backend. It has already been unmarshalled
  57. // from the main config file 'backend' config "backend_config"
  58. // Now we need to convert each type and copy into the guerrillaDBAndRedisConfig struct
  59. func (g *GuerrillaDBAndRedisBackend) loadConfig(backendConfig BackendConfig) (err error) {
  60. configType := BaseConfig(&guerrillaDBAndRedisConfig{})
  61. bcfg, err := Svc.ExtractConfig(backendConfig, configType)
  62. if err != nil {
  63. return err
  64. }
  65. m := bcfg.(*guerrillaDBAndRedisConfig)
  66. g.config = m
  67. return nil
  68. }
  69. func (g *GuerrillaDBAndRedisBackend) getNumberOfWorkers() int {
  70. return g.config.NumberOfWorkers
  71. }
  72. type redisClient struct {
  73. isConnected bool
  74. conn RedisConn
  75. time int
  76. }
  77. // compressedData struct will be compressed using zlib when printed via fmt
  78. type compressedData struct {
  79. extraHeaders []byte
  80. data *bytes.Buffer
  81. pool *sync.Pool
  82. }
  83. // newCompressedData returns a new CompressedData
  84. func newCompressedData() *compressedData {
  85. var p = sync.Pool{
  86. New: func() interface{} {
  87. var b bytes.Buffer
  88. return &b
  89. },
  90. }
  91. return &compressedData{
  92. pool: &p,
  93. }
  94. }
  95. // Set the extraheaders and buffer of data to compress
  96. func (c *compressedData) set(b []byte, d *bytes.Buffer) {
  97. c.extraHeaders = b
  98. c.data = d
  99. }
  100. // implement Stringer interface
  101. func (c *compressedData) String() string {
  102. if c.data == nil {
  103. return ""
  104. }
  105. //borrow a buffer form the pool
  106. b := c.pool.Get().(*bytes.Buffer)
  107. // put back in the pool
  108. defer func() {
  109. b.Reset()
  110. c.pool.Put(b)
  111. }()
  112. var r *bytes.Reader
  113. w, _ := zlib.NewWriterLevel(b, zlib.BestSpeed)
  114. r = bytes.NewReader(c.extraHeaders)
  115. _, _ = io.Copy(w, r)
  116. _, _ = io.Copy(w, c.data)
  117. _ = w.Close()
  118. return b.String()
  119. }
  120. // clear it, without clearing the pool
  121. func (c *compressedData) clear() {
  122. c.extraHeaders = []byte{}
  123. c.data = nil
  124. }
  125. // prepares the sql query with the number of rows that can be batched with it
  126. func (g *GuerrillaDBAndRedisBackend) prepareInsertQuery(rows int, db *sql.DB) *sql.Stmt {
  127. if rows == 0 {
  128. panic("rows argument cannot be 0")
  129. }
  130. if g.cache[rows-1] != nil {
  131. return g.cache[rows-1]
  132. }
  133. sqlstr := "INSERT INTO " + g.config.Table + "" +
  134. "(" +
  135. "`date`, " +
  136. "`to`, " +
  137. "`from`, " +
  138. "`subject`, " +
  139. "`body`, " +
  140. "`charset`, " +
  141. "`mail`, " +
  142. "`spam_score`, " +
  143. "`hash`, " +
  144. "`content_type`, " +
  145. "`recipient`, " +
  146. "`has_attach`, " +
  147. "`ip_addr`, " +
  148. "`return_path`, " +
  149. "`is_tls`" +
  150. ")" +
  151. " values "
  152. values := "(NOW(), ?, ?, ?, ? , 'UTF-8' , ?, 0, ?, '', ?, 0, ?, ?, ?)"
  153. // add more rows
  154. comma := ""
  155. for i := 0; i < rows; i++ {
  156. sqlstr += comma + values
  157. if comma == "" {
  158. comma = ","
  159. }
  160. }
  161. stmt, sqlErr := db.Prepare(sqlstr)
  162. if sqlErr != nil {
  163. Log().WithError(sqlErr).Fatalf("failed while db.Prepare(INSERT...)")
  164. }
  165. // cache it
  166. g.cache[rows-1] = stmt
  167. return stmt
  168. }
  169. func (g *GuerrillaDBAndRedisBackend) doQuery(c int, db *sql.DB, insertStmt *sql.Stmt, vals *[]interface{}) error {
  170. var execErr error
  171. defer func() {
  172. if r := recover(); r != nil {
  173. //logln(1, fmt.Sprintf("Recovered in %v", r))
  174. Log().Error("Recovered form panic:", r, string(debug.Stack()))
  175. sum := 0
  176. for _, v := range *vals {
  177. if str, ok := v.(string); ok {
  178. sum = sum + len(str)
  179. }
  180. }
  181. Log().Errorf("panic while inserting query [%s] size:%d, err %v", r, sum, execErr)
  182. panic("query failed")
  183. }
  184. }()
  185. // prepare the query used to insert when rows reaches batchMax
  186. insertStmt = g.prepareInsertQuery(c, db)
  187. _, execErr = insertStmt.Exec(*vals...)
  188. //if rand.Intn(2) == 1 {
  189. // return errors.New("uggabooka")
  190. //}
  191. if execErr != nil {
  192. Log().WithError(execErr).Error("There was a problem the insert")
  193. }
  194. return execErr
  195. }
  196. // Batches the rows from the feeder chan in to a single INSERT statement.
  197. // Execute the batches query when:
  198. // - number of batched rows reaches a threshold, i.e. count n = threshold
  199. // - or, no new rows within a certain time, i.e. times out
  200. // The goroutine can either exit if there's a panic or feeder channel closes
  201. // it returns feederOk which signals if the feeder chanel was ok (still open) while returning
  202. // if it feederOk is false, then it means the feeder chanel is closed
  203. func (g *GuerrillaDBAndRedisBackend) insertQueryBatcher(
  204. feeder feedChan,
  205. db *sql.DB,
  206. batcherId int,
  207. stop chan bool) (feederOk bool) {
  208. // controls shutdown
  209. defer g.batcherWg.Done()
  210. g.batcherWg.Add(1)
  211. // vals is where values are batched to
  212. var vals []interface{}
  213. // how many rows were batched
  214. count := 0
  215. // The timer will tick x seconds.
  216. // Interrupting the select clause when there's no data on the feeder channel
  217. timeo := GuerrillaDBAndRedisBatchTimeout
  218. if g.config.BatchTimeout > 0 {
  219. timeo = time.Duration(g.config.BatchTimeout)
  220. }
  221. t := time.NewTimer(timeo)
  222. // prepare the query used to insert when rows reaches batchMax
  223. insertStmt := g.prepareInsertQuery(GuerrillaDBAndRedisBatchMax, db)
  224. // inserts executes a batched insert query, clears the vals and resets the count
  225. inserter := func(c int) {
  226. if c > 0 {
  227. err := g.doQuery(c, db, insertStmt, &vals)
  228. if err != nil {
  229. // maybe connection prob?
  230. // retry the sql query
  231. attempts := 3
  232. for i := 0; i < attempts; i++ {
  233. Log().Infof("retrying query query rows[%c] ", c)
  234. time.Sleep(time.Second)
  235. err = g.doQuery(c, db, insertStmt, &vals)
  236. if err == nil {
  237. continue
  238. }
  239. }
  240. }
  241. }
  242. vals = nil
  243. count = 0
  244. }
  245. rand.Seed(time.Now().UnixNano())
  246. defer func() {
  247. if r := recover(); r != nil {
  248. Log().Error("insertQueryBatcher caught a panic", r, string(debug.Stack()))
  249. }
  250. }()
  251. // Keep getting values from feeder and add to batch.
  252. // if feeder times out, execute the batched query
  253. // otherwise, execute the batched query once it reaches the GuerrillaDBAndRedisBatchMax threshold
  254. feederOk = true
  255. for {
  256. select {
  257. // it may panic when reading on a closed feeder channel. feederOK detects if it was closed
  258. case <-stop:
  259. Log().Infof("MySQL query batcher stopped (#%d)", batcherId)
  260. // Insert any remaining rows
  261. inserter(count)
  262. feederOk = false
  263. close(feeder)
  264. return
  265. case row := <-feeder:
  266. vals = append(vals, row...)
  267. count++
  268. Log().Debug("new feeder row:", row, " cols:", len(row), " count:", count, " worker", batcherId)
  269. if count >= GuerrillaDBAndRedisBatchMax {
  270. inserter(GuerrillaDBAndRedisBatchMax)
  271. }
  272. // stop timer from firing (reset the interrupt)
  273. if !t.Stop() {
  274. // darin the timer
  275. <-t.C
  276. }
  277. t.Reset(timeo)
  278. case <-t.C:
  279. // anything to insert?
  280. if n := len(vals); n > 0 {
  281. inserter(count)
  282. }
  283. t.Reset(timeo)
  284. }
  285. }
  286. }
  287. func trimToLimit(str string, limit int) string {
  288. ret := strings.TrimSpace(str)
  289. if len(str) > limit {
  290. ret = str[:limit]
  291. }
  292. return ret
  293. }
  294. func (g *GuerrillaDBAndRedisBackend) sqlConnect() (*sql.DB, error) {
  295. if db, err := sql.Open(g.config.Driver, g.config.DSN); err != nil {
  296. Log().Error("cannot open database", err, "]")
  297. return nil, err
  298. } else {
  299. // do we have access?
  300. _, err = db.Query("SELECT mail_id FROM " + g.config.Table + " LIMIT 1")
  301. if err != nil {
  302. Log().Error("cannot select table:", err)
  303. return nil, err
  304. }
  305. return db, nil
  306. }
  307. }
  308. func (c *redisClient) redisConnection(redisInterface string) (err error) {
  309. if c.isConnected == false {
  310. c.conn, err = RedisDialer("tcp", redisInterface)
  311. if err != nil {
  312. // handle error
  313. return err
  314. }
  315. c.isConnected = true
  316. }
  317. return nil
  318. }
  319. type feedChan chan []interface{}
  320. // GuerrillaDbRedis is a specialized processor for Guerrilla mail. It is here as an example.
  321. // It's an example of a 'monolithic' processor.
  322. func GuerrillaDbRedis() Decorator {
  323. g := GuerrillaDBAndRedisBackend{}
  324. redisClient := &redisClient{}
  325. var (
  326. db *sql.DB
  327. to, body string
  328. redisErr error
  329. feeders []feedChan
  330. )
  331. g.batcherStoppers = make([]chan bool, 0)
  332. Svc.AddInitializer(InitializeWith(func(backendConfig BackendConfig) error {
  333. configType := BaseConfig(&guerrillaDBAndRedisConfig{})
  334. bcfg, err := Svc.ExtractConfig(backendConfig, configType)
  335. if err != nil {
  336. return err
  337. }
  338. g.config = bcfg.(*guerrillaDBAndRedisConfig)
  339. db, err = g.sqlConnect()
  340. if err != nil {
  341. return err
  342. }
  343. queryBatcherId++
  344. // start the query SQL batching where we will send data via the feeder channel
  345. stop := make(chan bool)
  346. feeder := make(feedChan, 1)
  347. go func(qbID int, stop chan bool) {
  348. // we loop so that if insertQueryBatcher panics, it can recover and go in again
  349. for {
  350. if feederOK := g.insertQueryBatcher(feeder, db, qbID, stop); !feederOK {
  351. Log().Debugf("insertQueryBatcher exited (#%d)", qbID)
  352. return
  353. }
  354. Log().Debug("resuming insertQueryBatcher")
  355. }
  356. }(queryBatcherId, stop)
  357. g.batcherStoppers = append(g.batcherStoppers, stop)
  358. feeders = append(feeders, feeder)
  359. return nil
  360. }))
  361. Svc.AddShutdowner(ShutdownWith(func() error {
  362. if err := db.Close(); err != nil {
  363. Log().WithError(err).Error("close mysql failed")
  364. } else {
  365. Log().Infof("closed mysql")
  366. }
  367. if redisClient.conn != nil {
  368. if err := redisClient.conn.Close(); err != nil {
  369. Log().WithError(err).Error("close redis failed")
  370. } else {
  371. Log().Infof("closed redis")
  372. }
  373. }
  374. // send a close signal to all query batchers to exit.
  375. for i := range g.batcherStoppers {
  376. g.batcherStoppers[i] <- true
  377. }
  378. g.batcherWg.Wait()
  379. return nil
  380. }))
  381. var vals []interface{}
  382. data := newCompressedData()
  383. return func(p Processor) Processor {
  384. return ProcessWith(func(e *mail.Envelope, task SelectTask) (Result, error) {
  385. if task == TaskSaveMail {
  386. Log().Debug("Got mail from chan,", e.RemoteIP)
  387. to = trimToLimit(strings.TrimSpace(e.RcptTo[0].User)+"@"+g.config.PrimaryHost, 255)
  388. e.Helo = trimToLimit(e.Helo, 255)
  389. e.RcptTo[0].Host = trimToLimit(e.RcptTo[0].Host, 255)
  390. ts := fmt.Sprintf("%d", time.Now().UnixNano())
  391. if err := e.ParseHeaders(); err != nil {
  392. Log().WithError(err).Error("failed to parse headers")
  393. }
  394. hash := MD5Hex(
  395. to,
  396. e.MailFrom.String(),
  397. e.Subject,
  398. ts)
  399. e.QueuedId = hash
  400. // Add extra headers
  401. protocol := "SMTP"
  402. if e.ESMTP {
  403. protocol = "E" + protocol
  404. }
  405. if e.TLS {
  406. protocol = protocol + "S"
  407. }
  408. var addHead string
  409. addHead += "Delivered-To: " + to + "\r\n"
  410. addHead += "Received: from " + e.RemoteIP + " ([" + e.RemoteIP + "])\r\n"
  411. addHead += " by " + e.RcptTo[0].Host + " with " + protocol + " id " + hash + "@" + e.RcptTo[0].Host + ";\r\n"
  412. addHead += " " + time.Now().Format(time.RFC1123Z) + "\r\n"
  413. // data will be compressed when printed, with addHead added to beginning
  414. data.set([]byte(addHead), &e.Data)
  415. body = "gzencode"
  416. // data will be written to redis - it implements the Stringer interface, redigo uses fmt to
  417. // print the data to redis.
  418. redisErr = redisClient.redisConnection(g.config.RedisInterface)
  419. if redisErr == nil {
  420. _, doErr := redisClient.conn.Do("SETEX", hash, g.config.RedisExpireSeconds, data)
  421. if doErr == nil {
  422. body = "redis" // the backend system will know to look in redis for the message data
  423. data.clear() // blank
  424. }
  425. } else {
  426. Log().WithError(redisErr).Warn("Error while connecting redis")
  427. }
  428. vals = []interface{}{} // clear the vals
  429. vals = append(vals,
  430. trimToLimit(to, 255),
  431. trimToLimit(e.MailFrom.String(), 255),
  432. trimToLimit(e.Subject, 255),
  433. body,
  434. data.String(),
  435. hash,
  436. trimToLimit(to, 255),
  437. e.RemoteIP,
  438. trimToLimit(e.MailFrom.String(), 255),
  439. e.TLS)
  440. // give the values to a random query batcher
  441. feeders[rand.Intn(len(feeders))] <- vals
  442. return p.Process(e, task)
  443. } else {
  444. return p.Process(e, task)
  445. }
  446. })
  447. }
  448. }