tx.go 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868
  1. package storage
  2. import (
  3. "fmt"
  4. "sort"
  5. "strings"
  6. "sync"
  7. )
  8. // ErrSerialization is returned when a transaction's optimistic validation fails
  9. // because a concurrent transaction committed a conflicting change.
  10. var ErrSerialization = fmt.Errorf("serialization failure: concurrent transaction modified the database")
  11. // Session is a buffered SQL transaction/session wrapper around a TableManager.
  12. // While in a transaction, writes are staged in memory, reads observe that
  13. // staged overlay, and COMMIT issues a single atomic compare-and-swap batch
  14. // write. Rollback simply discards the staged changes (no compensating durable
  15. // writes). A Session is used for both autocommit statements (where operations
  16. // delegate straight to the durable TableManager) and buffered transactions.
  17. type Session struct {
  18. schema *SchemaManager
  19. table *TableManager
  20. mu sync.Mutex
  21. inTx bool
  22. aborted bool
  23. // overlay holds staged writes keyed by lowercased table then data key.
  24. overlay map[string]map[string]*overlayEntry
  25. // log records mutations in order so savepoints can roll back.
  26. log []txMutation
  27. // reads records the durable LSN of every key the transaction observed
  28. // (0 means the key was observed absent). It becomes the compare set at
  29. // commit and also captures the base LSN of every written key.
  30. reads map[string]uint64
  31. // scanGens records the per-table generation captured by the first scan of
  32. // each table, validated at commit to detect phantoms. Indexed equality reads
  33. // use predicateGens so writes to other index values do not cause conflicts.
  34. scanGens map[string]uint64
  35. predicateGens map[string]predicateRead
  36. }
  37. type predicateRead struct {
  38. table string
  39. gen uint64
  40. }
  41. type overlayEntry struct {
  42. row Row
  43. absent bool
  44. }
  45. type txMutation struct {
  46. table string
  47. key string
  48. prev *overlayEntry
  49. }
  50. // NewSession creates a session wrapping the given schema and table managers.
  51. func NewSession(schema *SchemaManager, table *TableManager) *Session {
  52. return &Session{
  53. schema: schema,
  54. table: table,
  55. overlay: make(map[string]map[string]*overlayEntry),
  56. reads: make(map[string]uint64),
  57. scanGens: make(map[string]uint64),
  58. predicateGens: make(map[string]predicateRead),
  59. }
  60. }
  61. // Begin starts a buffered transaction.
  62. func (s *Session) Begin() error {
  63. s.mu.Lock()
  64. defer s.mu.Unlock()
  65. if s.inTx {
  66. return fmt.Errorf("cannot start a transaction within a transaction")
  67. }
  68. s.inTx = true
  69. s.aborted = false
  70. s.overlay = make(map[string]map[string]*overlayEntry)
  71. s.log = nil
  72. s.reads = make(map[string]uint64)
  73. s.scanGens = make(map[string]uint64)
  74. s.predicateGens = make(map[string]predicateRead)
  75. return nil
  76. }
  77. // InTx reports whether a transaction is in progress.
  78. func (s *Session) InTx() bool {
  79. s.mu.Lock()
  80. defer s.mu.Unlock()
  81. return s.inTx
  82. }
  83. // Abort marks the transaction as aborted without discarding state.
  84. func (s *Session) Abort() {
  85. s.mu.Lock()
  86. defer s.mu.Unlock()
  87. if s.inTx {
  88. s.aborted = true
  89. }
  90. }
  91. // Snapshot returns the current mutation-log position for a savepoint.
  92. func (s *Session) Snapshot() int {
  93. s.mu.Lock()
  94. defer s.mu.Unlock()
  95. return len(s.log)
  96. }
  97. // RollbackTo discards mutations after the given savepoint position.
  98. func (s *Session) RollbackTo(pos int) {
  99. s.mu.Lock()
  100. defer s.mu.Unlock()
  101. for i := len(s.log) - 1; i >= pos && i >= 0; i-- {
  102. m := s.log[i]
  103. if m.prev == nil {
  104. delete(s.overlay[m.table], m.key)
  105. if len(s.overlay[m.table]) == 0 {
  106. delete(s.overlay, m.table)
  107. }
  108. } else {
  109. if s.overlay[m.table] == nil {
  110. s.overlay[m.table] = make(map[string]*overlayEntry)
  111. }
  112. s.overlay[m.table][m.key] = m.prev
  113. }
  114. }
  115. if pos < len(s.log) {
  116. s.log = s.log[:pos]
  117. }
  118. s.aborted = false
  119. }
  120. // Rollback discards the transaction without writing anything durable.
  121. func (s *Session) Rollback() error {
  122. s.mu.Lock()
  123. defer s.mu.Unlock()
  124. if !s.inTx {
  125. return fmt.Errorf("cannot rollback: no transaction in progress")
  126. }
  127. s.resetLocked()
  128. return nil
  129. }
  130. // Commit validates and durably applies the staged transaction in one atomic
  131. // compare-and-swap batch write. It discards the commit LSN; callers that need a
  132. // ordering cursor use CommitWithLSN.
  133. func (s *Session) Commit() error {
  134. _, err := s.CommitWithLSN()
  135. return err
  136. }
  137. // CommitWithLSN behaves like Commit and additionally returns the LSN assigned
  138. // to the durable batch. A read-only or empty transaction returns zero because
  139. // no new durable state was published.
  140. func (s *Session) CommitWithLSN() (uint64, error) {
  141. s.mu.Lock()
  142. defer s.mu.Unlock()
  143. if !s.inTx {
  144. return 0, fmt.Errorf("cannot commit: no transaction in progress")
  145. }
  146. if s.aborted {
  147. s.resetLocked()
  148. return 0, fmt.Errorf("current transaction is aborted")
  149. }
  150. // Acquire the per-table gates, escalating any table whose UNIQUE index
  151. // appeared after the initial unlocked probe. See acquireCommitLocks.
  152. _, _, unlockAll, uniqueTables, err := s.acquireCommitLocks()
  153. if err != nil {
  154. s.resetLocked()
  155. return 0, err
  156. }
  157. defer unlockAll()
  158. // Validate scan generations for phantom detection.
  159. for t, gen := range s.scanGens {
  160. if s.table.generation(t) != gen {
  161. s.resetLocked()
  162. return 0, ErrSerialization
  163. }
  164. }
  165. for key, predicate := range s.predicateGens {
  166. if s.table.predicateGeneration(key) != predicate.gen {
  167. s.resetLocked()
  168. return 0, ErrSerialization
  169. }
  170. }
  171. // Validate the final overlay of every unique-indexed written table against
  172. // durable rows. This runs under the exclusive gate acquired above, so it
  173. // cannot race a concurrent writer. Duplicate unique values fail the commit
  174. // with a "UNIQUE constraint failed" error rather than ErrSerialization.
  175. for t := range uniqueTables {
  176. entries := s.overlay[t]
  177. pending := make([]Row, 0, len(entries))
  178. // Exclude every overlay data key: rows being written, updated, or deleted
  179. // are all replaced by this transaction, so their durable unique values
  180. // must not self-conflict with the staged state (e.g. delete a row and
  181. // re-insert the same primary key with a different rowid).
  182. excluded := make(map[string]bool, len(entries))
  183. for key, e := range entries {
  184. excluded[key] = true
  185. if !e.absent {
  186. pending = append(pending, e.row)
  187. }
  188. }
  189. if err := s.table.validateUniqueRows(t, pending, excluded); err != nil {
  190. s.resetLocked()
  191. return 0, err
  192. }
  193. }
  194. // Build the compare set from every observed key.
  195. checks := make([]CompareCheck, 0, len(s.reads))
  196. for key, lsn := range s.reads {
  197. checks = append(checks, CompareCheck{Key: []byte(key), LSN: lsn})
  198. }
  199. sort.Slice(checks, func(i, j int) bool { return string(checks[i].Key) < string(checks[j].Key) })
  200. // Build the batch ops from the staged overlay.
  201. ops := make([]BatchOp, 0)
  202. for _, entries := range s.overlay {
  203. for key, e := range entries {
  204. if e.absent {
  205. ops = append(ops, BatchOp{Op: batchDelete, Key: []byte(key)})
  206. } else {
  207. data, err := encodeRow(e.row)
  208. if err != nil {
  209. s.resetLocked()
  210. return 0, err
  211. }
  212. ops = append(ops, BatchOp{Op: batchPut, Key: []byte(key), Value: data})
  213. }
  214. }
  215. }
  216. sort.Slice(ops, func(i, j int) bool { return string(ops[i].Key) < string(ops[j].Key) })
  217. // A read-only transaction has nothing to write; commit trivially.
  218. if len(ops) == 0 {
  219. s.resetLocked()
  220. return 0, nil
  221. }
  222. var committed bool
  223. var commitLSN uint64
  224. err = s.table.pool.WithClient(func(c *KVClient) error {
  225. lsn, ok, err := c.CompareBatchWrite(checks, ops, nil)
  226. commitLSN = lsn
  227. committed = ok
  228. return err
  229. })
  230. if err != nil {
  231. s.resetLocked()
  232. return 0, err
  233. }
  234. if !committed {
  235. s.resetLocked()
  236. return 0, ErrSerialization
  237. }
  238. // Advance generations and invalidate derived caches for written tables.
  239. for t := range s.overlay {
  240. s.table.InvalidateCache(t)
  241. s.table.bumpIndexPredicateWildcard(t)
  242. s.table.bumpGeneration(t)
  243. }
  244. s.resetLocked()
  245. return commitLSN, nil
  246. }
  247. // acquireCommitLocks acquires the per-table gates needed to commit the staged
  248. // transaction, returning the held locks, whether each is exclusive, an unlock
  249. // function, and the set of written tables that carry a UNIQUE index. A written
  250. // table with a UNIQUE index is locked exclusively so its final overlay can be
  251. // validated against durable rows. The uniqueness probe is re-checked under the
  252. // held gates so a concurrent CREATE UNIQUE INDEX cannot slip in after the probe
  253. // and leave a duplicate unvalidated.
  254. func (s *Session) acquireCommitLocks() (locks []*sync.RWMutex, exclusive []bool, unlockAll func(), uniqueTables map[string]bool, err error) {
  255. unlock := func(ls []*sync.RWMutex, ex []bool) {
  256. for i := len(ls) - 1; i >= 0; i-- {
  257. if ex[i] {
  258. ls[i].Unlock()
  259. } else {
  260. ls[i].RUnlock()
  261. }
  262. }
  263. }
  264. for {
  265. // Tables read through scans or predicates need an exclusive validation
  266. // gate. Tables that are only written use the shared publication gate,
  267. // allowing disjoint optimistic commits to proceed concurrently while
  268. // still excluding scanner commits.
  269. affected := make(map[string]bool)
  270. for t := range s.overlay {
  271. affected[t] = false
  272. }
  273. for t := range s.scanGens {
  274. affected[t] = true
  275. }
  276. for _, predicate := range s.predicateGens {
  277. affected[predicate.table] = true
  278. }
  279. // Tables written in this transaction that carry a UNIQUE index need the
  280. // exclusive gate so their final overlay can be validated against durable
  281. // rows without racing another writer.
  282. uniqueTables = make(map[string]bool)
  283. for t := range s.overlay {
  284. uniq, err := s.table.hasUniqueIndex(t)
  285. if err != nil {
  286. return nil, nil, nil, nil, err
  287. }
  288. if uniq {
  289. uniqueTables[t] = true
  290. affected[t] = true
  291. }
  292. }
  293. tables := make([]string, 0, len(affected))
  294. for t := range affected {
  295. tables = append(tables, t)
  296. }
  297. sort.Strings(tables)
  298. locks = make([]*sync.RWMutex, len(tables))
  299. exclusive = make([]bool, len(tables))
  300. for i, t := range tables {
  301. locks[i] = s.table.tableLock(t)
  302. exclusive[i] = affected[t]
  303. }
  304. for i, lock := range locks {
  305. if exclusive[i] {
  306. lock.Lock()
  307. } else {
  308. lock.RLock()
  309. }
  310. }
  311. // Re-check under the held gates. A concurrent CREATE UNIQUE INDEX can
  312. // only have completed before we acquired the gate (it needs the exclusive
  313. // gate), so this probe is authoritative; escalate and retry if a written
  314. // table gained a unique index.
  315. retry := false
  316. for i, t := range tables {
  317. if exclusive[i] {
  318. continue
  319. }
  320. if _, written := s.overlay[t]; !written {
  321. continue
  322. }
  323. uniq, err := s.table.hasUniqueIndex(t)
  324. if err != nil {
  325. unlock(locks, exclusive)
  326. return nil, nil, nil, nil, err
  327. }
  328. if uniq {
  329. retry = true
  330. break
  331. }
  332. }
  333. if retry {
  334. unlock(locks, exclusive)
  335. continue
  336. }
  337. unlockAll = func() { unlock(locks, exclusive) }
  338. return locks, exclusive, unlockAll, uniqueTables, nil
  339. }
  340. }
  341. func (s *Session) resetLocked() {
  342. s.inTx = false
  343. s.aborted = false
  344. s.overlay = make(map[string]map[string]*overlayEntry)
  345. s.log = nil
  346. s.reads = make(map[string]uint64)
  347. s.scanGens = make(map[string]uint64)
  348. s.predicateGens = make(map[string]predicateRead)
  349. }
  350. func (s *Session) stagePut(table, key string, row Row) {
  351. tl := strings.ToLower(table)
  352. if s.overlay[tl] == nil {
  353. s.overlay[tl] = make(map[string]*overlayEntry)
  354. }
  355. s.log = append(s.log, txMutation{table: tl, key: key, prev: s.overlay[tl][key]})
  356. s.overlay[tl][key] = &overlayEntry{row: cloneRow(row)}
  357. if _, ok := s.reads[key]; !ok {
  358. s.reads[key] = 0
  359. }
  360. }
  361. func (s *Session) stageDelete(table, key string, row Row) {
  362. tl := strings.ToLower(table)
  363. if s.overlay[tl] == nil {
  364. s.overlay[tl] = make(map[string]*overlayEntry)
  365. }
  366. s.log = append(s.log, txMutation{table: tl, key: key, prev: s.overlay[tl][key]})
  367. // Keep the deleted row so COMMIT can exempt its rowid from the unique-index
  368. // scan and release its unique value.
  369. s.overlay[tl][key] = &overlayEntry{row: cloneRow(row), absent: true}
  370. if _, ok := s.reads[key]; !ok {
  371. s.reads[key] = 0
  372. }
  373. }
  374. // GetByPK retrieves a row by primary key, observing the staged overlay in a
  375. // transaction and recording the observed LSN for validation.
  376. func (s *Session) GetByPK(table, pk string) (Row, error) {
  377. s.mu.Lock()
  378. defer s.mu.Unlock()
  379. if !s.inTx {
  380. return s.table.GetByPK(table, pk)
  381. }
  382. key := s.table.dataKey(table, pk)
  383. tl := strings.ToLower(table)
  384. if e, ok := s.overlay[tl][key]; ok {
  385. if e.absent {
  386. return nil, ErrKeyNotFound
  387. }
  388. return cloneRow(e.row), nil
  389. }
  390. row, lsn, err := s.table.getByPKWithLSN(table, pk)
  391. if err != nil {
  392. if err == ErrKeyNotFound {
  393. s.reads[key] = 0
  394. return nil, ErrKeyNotFound
  395. }
  396. return nil, err
  397. }
  398. s.reads[key] = lsn
  399. return row, nil
  400. }
  401. // Select scans a table, merging the staged overlay so a transaction sees its
  402. // own writes, and records per-row LSNs plus the table generation.
  403. func (s *Session) Select(table string, filter func(Row) bool) ([]Row, error) {
  404. s.mu.Lock()
  405. if !s.inTx {
  406. s.mu.Unlock()
  407. return s.table.Select(table, filter)
  408. }
  409. defer s.mu.Unlock()
  410. return s.selectLocked(table, filter)
  411. }
  412. func (s *Session) selectLocked(table string, filter func(Row) bool) ([]Row, error) {
  413. schema, err := s.schema.GetSchema(table)
  414. if err != nil {
  415. return nil, err
  416. }
  417. tl := strings.ToLower(table)
  418. tableLock := s.table.tableLock(tl)
  419. tableLock.RLock()
  420. defer tableLock.RUnlock()
  421. if _, ok := s.scanGens[tl]; !ok {
  422. s.scanGens[tl] = s.table.generation(tl)
  423. }
  424. overlay := s.overlay[tl]
  425. var rows []Row
  426. err = s.table.scanRowsWithLSN(table, func(row Row, lsn uint64) (bool, error) {
  427. key := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey]))
  428. if _, ok := overlay[key]; ok {
  429. return false, nil
  430. }
  431. s.reads[key] = lsn
  432. if filter == nil || filter(row) {
  433. rows = append(rows, row)
  434. }
  435. return false, nil
  436. })
  437. if err != nil {
  438. return nil, err
  439. }
  440. for _, e := range overlay {
  441. if e.absent {
  442. continue
  443. }
  444. if filter == nil || filter(e.row) {
  445. rows = append(rows, cloneRow(e.row))
  446. }
  447. }
  448. return rows, nil
  449. }
  450. // SelectByIndex reads only matching durable rows while capturing their LSNs,
  451. // then merges the transaction overlay. The table generation protects against
  452. // matching rows being inserted or removed after the lookup.
  453. func (s *Session) SelectByIndex(table, indexName string, colValue interface{}) ([]Row, error) {
  454. s.mu.Lock()
  455. defer s.mu.Unlock()
  456. if !s.inTx {
  457. return s.table.SelectByIndex(table, indexName, colValue)
  458. }
  459. tableKey := strings.ToLower(table)
  460. index, err := s.schema.GetIndex(indexName)
  461. if err != nil {
  462. return nil, err
  463. }
  464. col := index.Columns[0].Name
  465. want := formatIndexValue(colValue)
  466. versions, predicate, err := s.table.selectByIndexWithLSN(table, indexName, colValue)
  467. if err != nil {
  468. return nil, err
  469. }
  470. if _, ok := s.predicateGens[predicate.valueKey]; !ok {
  471. s.predicateGens[predicate.valueKey] = predicateRead{table: tableKey, gen: predicate.valueGen}
  472. }
  473. if _, ok := s.predicateGens[predicate.wildcardKey]; !ok {
  474. s.predicateGens[predicate.wildcardKey] = predicateRead{table: tableKey, gen: predicate.wildcardGen}
  475. }
  476. overlay := s.overlay[tableKey]
  477. rows := make([]Row, 0, len(versions)+len(overlay))
  478. for _, version := range versions {
  479. if _, staged := overlay[version.key]; staged {
  480. continue
  481. }
  482. s.reads[version.key] = version.lsn
  483. rows = append(rows, version.row)
  484. }
  485. for _, entry := range overlay {
  486. if !entry.absent && formatIndexValue(entry.row[col]) == want {
  487. rows = append(rows, cloneRow(entry.row))
  488. }
  489. }
  490. return rows, nil
  491. }
  492. // SelectByIndexKey retrieves the rows whose index key equals valueKey, merging
  493. // the staged overlay so a transaction sees its own writes. valueKey is the
  494. // formatted key from TableManager.IndexRowKey, which supports plain, composite,
  495. // and expression indexes uniformly.
  496. func (s *Session) SelectByIndexKey(table string, idx *Index, valueKey string) ([]Row, error) {
  497. s.mu.Lock()
  498. defer s.mu.Unlock()
  499. if !s.inTx {
  500. versions, _, err := s.table.selectByIndexKeyWithLSN(table, idx.Name, valueKey)
  501. if err != nil {
  502. return nil, err
  503. }
  504. rows := make([]Row, len(versions))
  505. for i := range versions {
  506. rows[i] = versions[i].row
  507. }
  508. return rows, nil
  509. }
  510. tableKey := strings.ToLower(table)
  511. versions, predicate, err := s.table.selectByIndexKeyWithLSN(table, idx.Name, valueKey)
  512. if err != nil {
  513. return nil, err
  514. }
  515. if _, ok := s.predicateGens[predicate.valueKey]; !ok {
  516. s.predicateGens[predicate.valueKey] = predicateRead{table: tableKey, gen: predicate.valueGen}
  517. }
  518. if _, ok := s.predicateGens[predicate.wildcardKey]; !ok {
  519. s.predicateGens[predicate.wildcardKey] = predicateRead{table: tableKey, gen: predicate.wildcardGen}
  520. }
  521. overlay := s.overlay[tableKey]
  522. rows := make([]Row, 0, len(versions)+len(overlay))
  523. for _, version := range versions {
  524. if _, staged := overlay[version.key]; staged {
  525. continue
  526. }
  527. s.reads[version.key] = version.lsn
  528. rows = append(rows, version.row)
  529. }
  530. for _, entry := range overlay {
  531. if entry.absent {
  532. continue
  533. }
  534. key, kerr := s.table.IndexRowKey(idx, entry.row)
  535. if kerr != nil {
  536. return nil, kerr
  537. }
  538. if key == valueKey {
  539. rows = append(rows, cloneRow(entry.row))
  540. }
  541. }
  542. return rows, nil
  543. }
  544. // CountFast returns the exact row count, observing the staged overlay in a
  545. // transaction.
  546. func (s *Session) CountFast(table string) (int, error) {
  547. s.mu.Lock()
  548. defer s.mu.Unlock()
  549. if !s.inTx {
  550. return s.table.CountFast(table)
  551. }
  552. rows, err := s.selectLocked(table, nil)
  553. if err != nil {
  554. return 0, err
  555. }
  556. return len(rows), nil
  557. }
  558. // Insert stages an insert in a transaction, or performs a durable autocommit
  559. // insert otherwise.
  560. func (s *Session) Insert(table string, row Row) error {
  561. _, err := s.InsertWithRowID(table, row)
  562. return err
  563. }
  564. // InsertWithRowID stages an insert in a transaction, or performs a durable
  565. // autocommit insert otherwise, returning the actual generated ROWID.
  566. func (s *Session) InsertWithRowID(table string, row Row) (int64, error) {
  567. s.mu.Lock()
  568. defer s.mu.Unlock()
  569. if !s.inTx {
  570. return s.table.InsertWithRowID(table, row)
  571. }
  572. return s.insertLocked(table, row)
  573. }
  574. // InsertBulk stages or durably bulk-inserts multiple rows.
  575. func (s *Session) InsertBulk(table string, rows []Row) (int, error) {
  576. count, _, err := s.InsertBulkWithLastRowID(table, rows)
  577. return count, err
  578. }
  579. // InsertBulkWithLastRowID is InsertBulk, additionally returning the ROWID of the
  580. // last staged/persisted row (0 when nothing was inserted).
  581. func (s *Session) InsertBulkWithLastRowID(table string, rows []Row) (int, int64, error) {
  582. s.mu.Lock()
  583. defer s.mu.Unlock()
  584. if !s.inTx {
  585. return s.table.InsertBulkWithLastRowID(table, rows)
  586. }
  587. var lastRowID int64
  588. count := 0
  589. for _, row := range rows {
  590. rid, err := s.insertLocked(table, row)
  591. if err != nil {
  592. return count, lastRowID, err
  593. }
  594. lastRowID = rid
  595. count++
  596. }
  597. return count, lastRowID, nil
  598. }
  599. // insertLocked is the transaction insert helper (caller holds s.mu).
  600. func (s *Session) insertLocked(table string, row Row) (int64, error) {
  601. nr, key, err := s.table.prepareInsert(table, row)
  602. if err != nil {
  603. return 0, err
  604. }
  605. rowid, _ := rowIDFromRow(nr)
  606. schema, err := s.schema.GetSchema(table)
  607. if err != nil {
  608. return 0, err
  609. }
  610. pk := fmt.Sprintf("%v", nr[schema.PrimaryKey])
  611. tl := strings.ToLower(table)
  612. if e, ok := s.overlay[tl][key]; ok {
  613. if !e.absent {
  614. return 0, fmt.Errorf("duplicate primary key: %s", pk)
  615. }
  616. } else {
  617. _, lsn, err := s.table.getByPKWithLSN(table, pk)
  618. if err == nil {
  619. s.reads[key] = lsn
  620. return 0, fmt.Errorf("duplicate primary key: %s", pk)
  621. }
  622. if err != ErrKeyNotFound {
  623. return 0, err
  624. }
  625. s.reads[key] = 0
  626. }
  627. s.stagePut(table, key, nr)
  628. return rowid, nil
  629. }
  630. // UpdateByPK stages or durably applies a single-row update.
  631. func (s *Session) UpdateByPK(table, pk string, updateFn func(Row) (Row, error)) (Row, bool, error) {
  632. s.mu.Lock()
  633. if !s.inTx {
  634. s.mu.Unlock()
  635. return s.table.UpdateByPK(table, pk, updateFn)
  636. }
  637. defer s.mu.Unlock()
  638. schema, err := s.schema.GetSchema(table)
  639. if err != nil {
  640. return nil, false, err
  641. }
  642. key := s.table.dataKey(table, pk)
  643. row, err := s.getByPKLocked(table, pk)
  644. if err == ErrKeyNotFound {
  645. return nil, false, nil
  646. }
  647. if err != nil {
  648. return nil, false, err
  649. }
  650. oldRow := cloneRow(row)
  651. updates, err := s.runUpdateFn(updateFn, row)
  652. if err != nil {
  653. return nil, false, err
  654. }
  655. for name, value := range updates {
  656. for _, column := range schema.Columns {
  657. if strings.EqualFold(name, column.Name) {
  658. row[column.Name] = value
  659. break
  660. }
  661. }
  662. }
  663. s.stagePut(table, key, row)
  664. return oldRow, true, nil
  665. }
  666. // runUpdateFn releases the session lock while invoking the update callback so the
  667. // callback can run nested reads (e.g. a scalar subquery in SET) through the same
  668. // session without deadlocking on s.mu. A connection executes one statement at a
  669. // time, so the staged overlay cannot change while the callback runs. The lock is
  670. // re-acquired before returning.
  671. func (s *Session) runUpdateFn(updateFn func(Row) (Row, error), row Row) (Row, error) {
  672. s.mu.Unlock()
  673. defer s.mu.Lock()
  674. return updateFn(row)
  675. }
  676. // DeleteByPK stages or durably applies a single-row delete.
  677. func (s *Session) DeleteByPK(table, pk string) (Row, bool, error) {
  678. s.mu.Lock()
  679. defer s.mu.Unlock()
  680. if !s.inTx {
  681. return s.table.DeleteByPK(table, pk)
  682. }
  683. key := s.table.dataKey(table, pk)
  684. row, err := s.getByPKLocked(table, pk)
  685. if err == ErrKeyNotFound {
  686. return nil, false, nil
  687. }
  688. if err != nil {
  689. return nil, false, err
  690. }
  691. s.stageDelete(table, key, row)
  692. return row, true, nil
  693. }
  694. // getByPKLocked reads a row observing the overlay (caller holds s.mu).
  695. func (s *Session) getByPKLocked(table, pk string) (Row, error) {
  696. key := s.table.dataKey(table, pk)
  697. tl := strings.ToLower(table)
  698. if e, ok := s.overlay[tl][key]; ok {
  699. if e.absent {
  700. return nil, ErrKeyNotFound
  701. }
  702. return cloneRow(e.row), nil
  703. }
  704. row, lsn, err := s.table.getByPKWithLSN(table, pk)
  705. if err != nil {
  706. if err == ErrKeyNotFound {
  707. s.reads[key] = 0
  708. return nil, ErrKeyNotFound
  709. }
  710. return nil, err
  711. }
  712. s.reads[key] = lsn
  713. return row, nil
  714. }
  715. // UpdateFunc stages or durably applies a scan-based update.
  716. func (s *Session) UpdateFunc(table string, updateFn func(Row) (Row, error), filter func(Row) bool) (int, error) {
  717. s.mu.Lock()
  718. if !s.inTx {
  719. s.mu.Unlock()
  720. return s.table.UpdateFunc(table, updateFn, filter)
  721. }
  722. defer s.mu.Unlock()
  723. schema, err := s.schema.GetSchema(table)
  724. if err != nil {
  725. return 0, err
  726. }
  727. rows, err := s.selectLocked(table, filter)
  728. if err != nil {
  729. return 0, err
  730. }
  731. count := 0
  732. for _, row := range rows {
  733. oldRow := cloneRow(row)
  734. oldKey := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey]))
  735. updates, err := s.runUpdateFn(updateFn, row)
  736. if err != nil {
  737. return count, err
  738. }
  739. for name, value := range updates {
  740. for _, column := range schema.Columns {
  741. if strings.EqualFold(name, column.Name) {
  742. row[column.Name] = value
  743. break
  744. }
  745. }
  746. }
  747. newKey := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey]))
  748. if newKey != oldKey {
  749. // The primary key changed: move the staged row instead of leaving a
  750. // stale copy under the old key, and refuse to overwrite another row.
  751. if _, gerr := s.getByPKLocked(table, fmt.Sprintf("%v", row[schema.PrimaryKey])); gerr == nil {
  752. return count, fmt.Errorf("duplicate primary key: %v", row[schema.PrimaryKey])
  753. } else if gerr != ErrKeyNotFound {
  754. return count, gerr
  755. }
  756. s.stageDelete(table, oldKey, oldRow)
  757. }
  758. s.stagePut(table, newKey, row)
  759. count++
  760. }
  761. return count, nil
  762. }
  763. // Delete stages or durably applies a scan-based delete.
  764. func (s *Session) Delete(table string, filter func(Row) bool) (int, error) {
  765. s.mu.Lock()
  766. if !s.inTx {
  767. s.mu.Unlock()
  768. return s.table.Delete(table, filter)
  769. }
  770. defer s.mu.Unlock()
  771. schema, err := s.schema.GetSchema(table)
  772. if err != nil {
  773. return 0, err
  774. }
  775. rows, err := s.selectLocked(table, filter)
  776. if err != nil {
  777. return 0, err
  778. }
  779. count := 0
  780. for _, row := range rows {
  781. key := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey]))
  782. s.stageDelete(table, key, row)
  783. count++
  784. }
  785. return count, nil
  786. }
  787. // ClearIndex passes through to the durable table manager.
  788. func (s *Session) ClearIndex(indexName, tableName string, columns []string) error {
  789. return s.table.ClearIndex(indexName, tableName, columns)
  790. }
  791. // InvalidateCache passes through to the durable table manager.
  792. func (s *Session) InvalidateCache(table string) {
  793. s.table.InvalidateCache(table)
  794. }
  795. // BuildIndex passes through to the durable table manager.
  796. func (s *Session) BuildIndex(indexName, tableName string, columns []string) error {
  797. return s.table.BuildIndex(indexName, tableName, columns)
  798. }