/go/raft/rel_store.go

https://github.com/github/orchestrator · Go · 249 lines · 199 code · 24 blank · 26 comment · 49 complexity · 98186ac7d86752ce89a667c19408dc20 MD5 · raw file

  1. /*
  2. Copyright 2017 Shlomi Noach, GitHub Inc.
  3. Licensed under the Apache License, Version 2.0 (the "License");
  4. you may not use this file except in compliance with the License.
  5. You may obtain a copy of the License at
  6. http://www.apache.org/licenses/LICENSE-2.0
  7. Unless required by applicable law or agreed to in writing, software
  8. distributed under the License is distributed on an "AS IS" BASIS,
  9. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  10. See the License for the specific language governing permissions and
  11. limitations under the License.
  12. */
  13. package orcraft
  14. import (
  15. "database/sql"
  16. "encoding/binary"
  17. "path/filepath"
  18. "sync"
  19. "github.com/openark/golib/log"
  20. "github.com/openark/golib/sqlutils"
  21. "github.com/hashicorp/raft"
  22. )
  23. const raftStoreFile = "raft_store.db"
  24. var createQueries = []string{
  25. `
  26. CREATE TABLE IF NOT EXISTS raft_log (
  27. log_index integer,
  28. term bigint not null,
  29. log_type int not null,
  30. data blob not null,
  31. PRIMARY KEY (log_index)
  32. )
  33. `,
  34. `
  35. CREATE TABLE IF NOT EXISTS raft_store (
  36. store_id integer,
  37. store_key varbinary(512) not null,
  38. store_value blob not null,
  39. PRIMARY KEY (store_id)
  40. )
  41. `,
  42. `
  43. CREATE INDEX IF NOT EXISTS store_key_idx_raft_store ON raft_store (store_key)
  44. `,
  45. }
  46. var dbMutex sync.Mutex
  47. // RelationalStoreimplements:
  48. // - hashicorp/raft.StableStore
  49. // - hashicorp/log.LogStore
  50. type RelationalStore struct {
  51. dataDir string
  52. backend *sql.DB
  53. }
  54. func NewRelationalStore(dataDir string) *RelationalStore {
  55. return &RelationalStore{
  56. dataDir: dataDir,
  57. }
  58. }
  59. func (relStore *RelationalStore) openDB() (*sql.DB, error) {
  60. dbMutex.Lock()
  61. defer dbMutex.Unlock()
  62. if relStore.backend == nil {
  63. relStoreFile := filepath.Join(relStore.dataDir, raftStoreFile)
  64. sqliteDB, _, err := sqlutils.GetSQLiteDB(relStoreFile)
  65. if err != nil {
  66. return nil, err
  67. }
  68. sqliteDB.SetMaxOpenConns(1)
  69. sqliteDB.SetMaxIdleConns(1)
  70. for _, query := range createQueries {
  71. if _, err := sqliteDB.Exec(sqlutils.ToSqlite3Dialect(query)); err != nil {
  72. return nil, err
  73. }
  74. }
  75. relStore.backend = sqliteDB
  76. log.Infof("raft: store initialized at %+v", relStoreFile)
  77. }
  78. return relStore.backend, nil
  79. }
  80. func (relStore *RelationalStore) Set(key []byte, val []byte) error {
  81. db, err := relStore.openDB()
  82. if err != nil {
  83. return err
  84. }
  85. tx, err := db.Begin()
  86. if err != nil {
  87. return err
  88. }
  89. stmt, err := tx.Prepare(`delete from raft_store where store_key = ?`)
  90. if err != nil {
  91. return err
  92. }
  93. _, err = stmt.Exec(key)
  94. if err != nil {
  95. tx.Rollback()
  96. return err
  97. }
  98. stmt, err = tx.Prepare(`insert into raft_store (store_key, store_value) values (?, ?)`)
  99. if err != nil {
  100. tx.Rollback()
  101. return err
  102. }
  103. _, err = stmt.Exec(key, val)
  104. if err != nil {
  105. tx.Rollback()
  106. return err
  107. }
  108. err = tx.Commit()
  109. return err
  110. }
  111. // Get returns the value for key, or an empty byte slice if key was not found.
  112. func (relStore *RelationalStore) Get(key []byte) (val []byte, err error) {
  113. db, err := relStore.openDB()
  114. if err != nil {
  115. return val, err
  116. }
  117. err = db.QueryRow("select min(store_value) from raft_store where store_key = ?", key).Scan(&val)
  118. return val, err
  119. }
  120. func (relStore *RelationalStore) SetUint64(key []byte, val uint64) error {
  121. b := make([]byte, 8)
  122. binary.LittleEndian.PutUint64(b, val)
  123. return relStore.Set(key, b)
  124. }
  125. // GetUint64 returns the uint64 value for key, or 0 if key was not found.
  126. func (relStore *RelationalStore) GetUint64(key []byte) (uint64, error) {
  127. b, err := relStore.Get(key)
  128. if err != nil {
  129. return 0, err
  130. }
  131. if len(b) == 0 {
  132. // Not found
  133. return 0, nil
  134. }
  135. i := binary.LittleEndian.Uint64(b)
  136. return i, nil
  137. }
  138. func (relStore *RelationalStore) FirstIndex() (idx uint64, err error) {
  139. db, err := relStore.openDB()
  140. if err != nil {
  141. return idx, err
  142. }
  143. err = db.QueryRow("select ifnull(min(log_index), 0) from raft_log").Scan(&idx)
  144. return idx, err
  145. }
  146. // LastIndex returns the last index written. 0 for no entries.
  147. func (relStore *RelationalStore) LastIndex() (idx uint64, err error) {
  148. db, err := relStore.openDB()
  149. if err != nil {
  150. return idx, err
  151. }
  152. err = db.QueryRow("select ifnull(max(log_index), 0) from raft_log").Scan(&idx)
  153. return idx, err
  154. }
  155. // GetLog gets a log entry at a given index.
  156. func (relStore *RelationalStore) GetLog(index uint64, log *raft.Log) error {
  157. db, err := relStore.openDB()
  158. if err != nil {
  159. return err
  160. }
  161. err = db.QueryRow(`
  162. select log_index, term, log_type, data
  163. from raft_log
  164. where log_index = ?
  165. `, index).Scan(&log.Index, &log.Term, &log.Type, &log.Data)
  166. if err == sql.ErrNoRows {
  167. return raft.ErrLogNotFound
  168. }
  169. return err
  170. }
  171. // StoreLog stores a log entry.
  172. func (relStore *RelationalStore) StoreLog(log *raft.Log) error {
  173. return relStore.StoreLogs([]*raft.Log{log})
  174. }
  175. // StoreLogs stores multiple log entries.
  176. func (relStore *RelationalStore) StoreLogs(logs []*raft.Log) error {
  177. db, err := relStore.openDB()
  178. if err != nil {
  179. return err
  180. }
  181. tx, err := db.Begin()
  182. if err != nil {
  183. return err
  184. }
  185. stmt, err := tx.Prepare(`
  186. replace into raft_log (
  187. log_index, term, log_type, data
  188. ) values (
  189. ?, ?, ?, ?
  190. )`)
  191. if err != nil {
  192. return err
  193. }
  194. for _, raftLog := range logs {
  195. _, err = stmt.Exec(raftLog.Index, raftLog.Term, int(raftLog.Type), raftLog.Data)
  196. if err != nil {
  197. tx.Rollback()
  198. return err
  199. }
  200. }
  201. return tx.Commit()
  202. }
  203. // DeleteRange deletes a range of log entries. The range is inclusive.
  204. func (relStore *RelationalStore) DeleteRange(min, max uint64) error {
  205. db, err := relStore.openDB()
  206. if err != nil {
  207. return err
  208. }
  209. _, err = db.Exec("delete from raft_log where log_index >= ? and log_index <= ?", min, max)
  210. return err
  211. }
  212. func (relStore *RelationalStore) DeleteAll() error {
  213. firstIndex, err := relStore.FirstIndex()
  214. if err != nil {
  215. return err
  216. }
  217. lastIndex, err := relStore.LastIndex()
  218. if err != nil {
  219. return err
  220. }
  221. return relStore.DeleteRange(firstIndex, lastIndex)
  222. }