package storage import ( "fmt" "sort" "strings" "sync" ) // ErrSerialization is returned when a transaction's optimistic validation fails // because a concurrent transaction committed a conflicting change. var ErrSerialization = fmt.Errorf("serialization failure: concurrent transaction modified the database") // Session is a buffered SQL transaction/session wrapper around a TableManager. // While in a transaction, writes are staged in memory, reads observe that // staged overlay, and COMMIT issues a single atomic compare-and-swap batch // write. Rollback simply discards the staged changes (no compensating durable // writes). A Session is used for both autocommit statements (where operations // delegate straight to the durable TableManager) and buffered transactions. type Session struct { schema *SchemaManager table *TableManager mu sync.Mutex inTx bool aborted bool // overlay holds staged writes keyed by lowercased table then data key. overlay map[string]map[string]*overlayEntry // log records mutations in order so savepoints can roll back. log []txMutation // reads records the durable LSN of every key the transaction observed // (0 means the key was observed absent). It becomes the compare set at // commit and also captures the base LSN of every written key. reads map[string]uint64 // scanGens records the per-table generation captured by the first scan of // each table, validated at commit to detect phantoms. Indexed equality reads // use predicateGens so writes to other index values do not cause conflicts. scanGens map[string]uint64 predicateGens map[string]predicateRead } type predicateRead struct { table string gen uint64 } type overlayEntry struct { row Row absent bool } type txMutation struct { table string key string prev *overlayEntry } // NewSession creates a session wrapping the given schema and table managers. func NewSession(schema *SchemaManager, table *TableManager) *Session { return &Session{ schema: schema, table: table, overlay: make(map[string]map[string]*overlayEntry), reads: make(map[string]uint64), scanGens: make(map[string]uint64), predicateGens: make(map[string]predicateRead), } } // Begin starts a buffered transaction. func (s *Session) Begin() error { s.mu.Lock() defer s.mu.Unlock() if s.inTx { return fmt.Errorf("cannot start a transaction within a transaction") } s.inTx = true s.aborted = false s.overlay = make(map[string]map[string]*overlayEntry) s.log = nil s.reads = make(map[string]uint64) s.scanGens = make(map[string]uint64) s.predicateGens = make(map[string]predicateRead) return nil } // InTx reports whether a transaction is in progress. func (s *Session) InTx() bool { s.mu.Lock() defer s.mu.Unlock() return s.inTx } // Abort marks the transaction as aborted without discarding state. func (s *Session) Abort() { s.mu.Lock() defer s.mu.Unlock() if s.inTx { s.aborted = true } } // Snapshot returns the current mutation-log position for a savepoint. func (s *Session) Snapshot() int { s.mu.Lock() defer s.mu.Unlock() return len(s.log) } // RollbackTo discards mutations after the given savepoint position. func (s *Session) RollbackTo(pos int) { s.mu.Lock() defer s.mu.Unlock() for i := len(s.log) - 1; i >= pos && i >= 0; i-- { m := s.log[i] if m.prev == nil { delete(s.overlay[m.table], m.key) if len(s.overlay[m.table]) == 0 { delete(s.overlay, m.table) } } else { if s.overlay[m.table] == nil { s.overlay[m.table] = make(map[string]*overlayEntry) } s.overlay[m.table][m.key] = m.prev } } if pos < len(s.log) { s.log = s.log[:pos] } s.aborted = false } // Rollback discards the transaction without writing anything durable. func (s *Session) Rollback() error { s.mu.Lock() defer s.mu.Unlock() if !s.inTx { return fmt.Errorf("cannot rollback: no transaction in progress") } s.resetLocked() return nil } // Commit validates and durably applies the staged transaction in one atomic // compare-and-swap batch write. func (s *Session) Commit() error { s.mu.Lock() defer s.mu.Unlock() if !s.inTx { return fmt.Errorf("cannot commit: no transaction in progress") } if s.aborted { s.resetLocked() return fmt.Errorf("current transaction is aborted") } // Acquire the per-table gates, escalating any table whose UNIQUE index // appeared after the initial unlocked probe. See acquireCommitLocks. _, _, unlockAll, uniqueTables, err := s.acquireCommitLocks() if err != nil { s.resetLocked() return err } defer unlockAll() // Validate scan generations for phantom detection. for t, gen := range s.scanGens { if s.table.generation(t) != gen { s.resetLocked() return ErrSerialization } } for key, predicate := range s.predicateGens { if s.table.predicateGeneration(key) != predicate.gen { s.resetLocked() return ErrSerialization } } // Validate the final overlay of every unique-indexed written table against // durable rows. This runs under the exclusive gate acquired above, so it // cannot race a concurrent writer. Duplicate unique values fail the commit // with a "UNIQUE constraint failed" error rather than ErrSerialization. for t := range uniqueTables { entries := s.overlay[t] pending := make([]Row, 0, len(entries)) // Exclude every overlay data key: rows being written, updated, or deleted // are all replaced by this transaction, so their durable unique values // must not self-conflict with the staged state (e.g. delete a row and // re-insert the same primary key with a different rowid). excluded := make(map[string]bool, len(entries)) for key, e := range entries { excluded[key] = true if !e.absent { pending = append(pending, e.row) } } if err := s.table.validateUniqueRows(t, pending, excluded); err != nil { s.resetLocked() return err } } // Build the compare set from every observed key. checks := make([]CompareCheck, 0, len(s.reads)) for key, lsn := range s.reads { checks = append(checks, CompareCheck{Key: []byte(key), LSN: lsn}) } sort.Slice(checks, func(i, j int) bool { return string(checks[i].Key) < string(checks[j].Key) }) // Build the batch ops from the staged overlay. ops := make([]BatchOp, 0) for _, entries := range s.overlay { for key, e := range entries { if e.absent { ops = append(ops, BatchOp{Op: batchDelete, Key: []byte(key)}) } else { data, err := encodeRow(e.row) if err != nil { s.resetLocked() return err } ops = append(ops, BatchOp{Op: batchPut, Key: []byte(key), Value: data}) } } } sort.Slice(ops, func(i, j int) bool { return string(ops[i].Key) < string(ops[j].Key) }) // A read-only transaction has nothing to write; commit trivially. if len(ops) == 0 { s.resetLocked() return nil } var committed bool err = s.table.pool.WithClient(func(c *KVClient) error { _, ok, err := c.CompareBatchWrite(checks, ops, nil) committed = ok return err }) if err != nil { s.resetLocked() return err } if !committed { s.resetLocked() return ErrSerialization } // Advance generations and invalidate derived caches for written tables. for t := range s.overlay { s.table.InvalidateCache(t) s.table.bumpIndexPredicateWildcard(t) s.table.bumpGeneration(t) } s.resetLocked() return nil } // acquireCommitLocks acquires the per-table gates needed to commit the staged // transaction, returning the held locks, whether each is exclusive, an unlock // function, and the set of written tables that carry a UNIQUE index. A written // table with a UNIQUE index is locked exclusively so its final overlay can be // validated against durable rows. The uniqueness probe is re-checked under the // held gates so a concurrent CREATE UNIQUE INDEX cannot slip in after the probe // and leave a duplicate unvalidated. func (s *Session) acquireCommitLocks() (locks []*sync.RWMutex, exclusive []bool, unlockAll func(), uniqueTables map[string]bool, err error) { unlock := func(ls []*sync.RWMutex, ex []bool) { for i := len(ls) - 1; i >= 0; i-- { if ex[i] { ls[i].Unlock() } else { ls[i].RUnlock() } } } for { // Tables read through scans or predicates need an exclusive validation // gate. Tables that are only written use the shared publication gate, // allowing disjoint optimistic commits to proceed concurrently while // still excluding scanner commits. affected := make(map[string]bool) for t := range s.overlay { affected[t] = false } for t := range s.scanGens { affected[t] = true } for _, predicate := range s.predicateGens { affected[predicate.table] = true } // Tables written in this transaction that carry a UNIQUE index need the // exclusive gate so their final overlay can be validated against durable // rows without racing another writer. uniqueTables = make(map[string]bool) for t := range s.overlay { uniq, err := s.table.hasUniqueIndex(t) if err != nil { return nil, nil, nil, nil, err } if uniq { uniqueTables[t] = true affected[t] = true } } tables := make([]string, 0, len(affected)) for t := range affected { tables = append(tables, t) } sort.Strings(tables) locks = make([]*sync.RWMutex, len(tables)) exclusive = make([]bool, len(tables)) for i, t := range tables { locks[i] = s.table.tableLock(t) exclusive[i] = affected[t] } for i, lock := range locks { if exclusive[i] { lock.Lock() } else { lock.RLock() } } // Re-check under the held gates. A concurrent CREATE UNIQUE INDEX can // only have completed before we acquired the gate (it needs the exclusive // gate), so this probe is authoritative; escalate and retry if a written // table gained a unique index. retry := false for i, t := range tables { if exclusive[i] { continue } if _, written := s.overlay[t]; !written { continue } uniq, err := s.table.hasUniqueIndex(t) if err != nil { unlock(locks, exclusive) return nil, nil, nil, nil, err } if uniq { retry = true break } } if retry { unlock(locks, exclusive) continue } unlockAll = func() { unlock(locks, exclusive) } return locks, exclusive, unlockAll, uniqueTables, nil } } func (s *Session) resetLocked() { s.inTx = false s.aborted = false s.overlay = make(map[string]map[string]*overlayEntry) s.log = nil s.reads = make(map[string]uint64) s.scanGens = make(map[string]uint64) s.predicateGens = make(map[string]predicateRead) } func (s *Session) stagePut(table, key string, row Row) { tl := strings.ToLower(table) if s.overlay[tl] == nil { s.overlay[tl] = make(map[string]*overlayEntry) } s.log = append(s.log, txMutation{table: tl, key: key, prev: s.overlay[tl][key]}) s.overlay[tl][key] = &overlayEntry{row: cloneRow(row)} if _, ok := s.reads[key]; !ok { s.reads[key] = 0 } } func (s *Session) stageDelete(table, key string, row Row) { tl := strings.ToLower(table) if s.overlay[tl] == nil { s.overlay[tl] = make(map[string]*overlayEntry) } s.log = append(s.log, txMutation{table: tl, key: key, prev: s.overlay[tl][key]}) // Keep the deleted row so COMMIT can exempt its rowid from the unique-index // scan and release its unique value. s.overlay[tl][key] = &overlayEntry{row: cloneRow(row), absent: true} if _, ok := s.reads[key]; !ok { s.reads[key] = 0 } } // GetByPK retrieves a row by primary key, observing the staged overlay in a // transaction and recording the observed LSN for validation. func (s *Session) GetByPK(table, pk string) (Row, error) { s.mu.Lock() defer s.mu.Unlock() if !s.inTx { return s.table.GetByPK(table, pk) } key := s.table.dataKey(table, pk) tl := strings.ToLower(table) if e, ok := s.overlay[tl][key]; ok { if e.absent { return nil, ErrKeyNotFound } return cloneRow(e.row), nil } row, lsn, err := s.table.getByPKWithLSN(table, pk) if err != nil { if err == ErrKeyNotFound { s.reads[key] = 0 return nil, ErrKeyNotFound } return nil, err } s.reads[key] = lsn return row, nil } // Select scans a table, merging the staged overlay so a transaction sees its // own writes, and records per-row LSNs plus the table generation. func (s *Session) Select(table string, filter func(Row) bool) ([]Row, error) { s.mu.Lock() if !s.inTx { s.mu.Unlock() return s.table.Select(table, filter) } defer s.mu.Unlock() return s.selectLocked(table, filter) } func (s *Session) selectLocked(table string, filter func(Row) bool) ([]Row, error) { schema, err := s.schema.GetSchema(table) if err != nil { return nil, err } tl := strings.ToLower(table) tableLock := s.table.tableLock(tl) tableLock.RLock() defer tableLock.RUnlock() if _, ok := s.scanGens[tl]; !ok { s.scanGens[tl] = s.table.generation(tl) } overlay := s.overlay[tl] var rows []Row err = s.table.scanRowsWithLSN(table, func(row Row, lsn uint64) (bool, error) { key := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey])) if _, ok := overlay[key]; ok { return false, nil } s.reads[key] = lsn if filter == nil || filter(row) { rows = append(rows, row) } return false, nil }) if err != nil { return nil, err } for _, e := range overlay { if e.absent { continue } if filter == nil || filter(e.row) { rows = append(rows, cloneRow(e.row)) } } return rows, nil } // SelectByIndex reads only matching durable rows while capturing their LSNs, // then merges the transaction overlay. The table generation protects against // matching rows being inserted or removed after the lookup. func (s *Session) SelectByIndex(table, indexName string, colValue interface{}) ([]Row, error) { s.mu.Lock() defer s.mu.Unlock() if !s.inTx { return s.table.SelectByIndex(table, indexName, colValue) } tableKey := strings.ToLower(table) index, err := s.schema.GetIndex(indexName) if err != nil { return nil, err } col := index.Columns[0].Name want := formatIndexValue(colValue) versions, predicate, err := s.table.selectByIndexWithLSN(table, indexName, colValue) if err != nil { return nil, err } if _, ok := s.predicateGens[predicate.valueKey]; !ok { s.predicateGens[predicate.valueKey] = predicateRead{table: tableKey, gen: predicate.valueGen} } if _, ok := s.predicateGens[predicate.wildcardKey]; !ok { s.predicateGens[predicate.wildcardKey] = predicateRead{table: tableKey, gen: predicate.wildcardGen} } overlay := s.overlay[tableKey] rows := make([]Row, 0, len(versions)+len(overlay)) for _, version := range versions { if _, staged := overlay[version.key]; staged { continue } s.reads[version.key] = version.lsn rows = append(rows, version.row) } for _, entry := range overlay { if !entry.absent && formatIndexValue(entry.row[col]) == want { rows = append(rows, cloneRow(entry.row)) } } return rows, nil } // SelectByIndexKey retrieves the rows whose index key equals valueKey, merging // the staged overlay so a transaction sees its own writes. valueKey is the // formatted key from TableManager.IndexRowKey, which supports plain, composite, // and expression indexes uniformly. func (s *Session) SelectByIndexKey(table string, idx *Index, valueKey string) ([]Row, error) { s.mu.Lock() defer s.mu.Unlock() if !s.inTx { versions, _, err := s.table.selectByIndexKeyWithLSN(table, idx.Name, valueKey) if err != nil { return nil, err } rows := make([]Row, len(versions)) for i := range versions { rows[i] = versions[i].row } return rows, nil } tableKey := strings.ToLower(table) versions, predicate, err := s.table.selectByIndexKeyWithLSN(table, idx.Name, valueKey) if err != nil { return nil, err } if _, ok := s.predicateGens[predicate.valueKey]; !ok { s.predicateGens[predicate.valueKey] = predicateRead{table: tableKey, gen: predicate.valueGen} } if _, ok := s.predicateGens[predicate.wildcardKey]; !ok { s.predicateGens[predicate.wildcardKey] = predicateRead{table: tableKey, gen: predicate.wildcardGen} } overlay := s.overlay[tableKey] rows := make([]Row, 0, len(versions)+len(overlay)) for _, version := range versions { if _, staged := overlay[version.key]; staged { continue } s.reads[version.key] = version.lsn rows = append(rows, version.row) } for _, entry := range overlay { if entry.absent { continue } key, kerr := s.table.IndexRowKey(idx, entry.row) if kerr != nil { return nil, kerr } if key == valueKey { rows = append(rows, cloneRow(entry.row)) } } return rows, nil } // CountFast returns the exact row count, observing the staged overlay in a // transaction. func (s *Session) CountFast(table string) (int, error) { s.mu.Lock() defer s.mu.Unlock() if !s.inTx { return s.table.CountFast(table) } rows, err := s.selectLocked(table, nil) if err != nil { return 0, err } return len(rows), nil } // Insert stages an insert in a transaction, or performs a durable autocommit // insert otherwise. func (s *Session) Insert(table string, row Row) error { _, err := s.InsertWithRowID(table, row) return err } // InsertWithRowID stages an insert in a transaction, or performs a durable // autocommit insert otherwise, returning the actual generated ROWID. func (s *Session) InsertWithRowID(table string, row Row) (int64, error) { s.mu.Lock() defer s.mu.Unlock() if !s.inTx { return s.table.InsertWithRowID(table, row) } return s.insertLocked(table, row) } // InsertBulk stages or durably bulk-inserts multiple rows. func (s *Session) InsertBulk(table string, rows []Row) (int, error) { count, _, err := s.InsertBulkWithLastRowID(table, rows) return count, err } // InsertBulkWithLastRowID is InsertBulk, additionally returning the ROWID of the // last staged/persisted row (0 when nothing was inserted). func (s *Session) InsertBulkWithLastRowID(table string, rows []Row) (int, int64, error) { s.mu.Lock() defer s.mu.Unlock() if !s.inTx { return s.table.InsertBulkWithLastRowID(table, rows) } var lastRowID int64 count := 0 for _, row := range rows { rid, err := s.insertLocked(table, row) if err != nil { return count, lastRowID, err } lastRowID = rid count++ } return count, lastRowID, nil } // insertLocked is the transaction insert helper (caller holds s.mu). func (s *Session) insertLocked(table string, row Row) (int64, error) { nr, key, err := s.table.prepareInsert(table, row) if err != nil { return 0, err } rowid, _ := rowIDFromRow(nr) schema, err := s.schema.GetSchema(table) if err != nil { return 0, err } pk := fmt.Sprintf("%v", nr[schema.PrimaryKey]) tl := strings.ToLower(table) if e, ok := s.overlay[tl][key]; ok { if !e.absent { return 0, fmt.Errorf("duplicate primary key: %s", pk) } } else { _, lsn, err := s.table.getByPKWithLSN(table, pk) if err == nil { s.reads[key] = lsn return 0, fmt.Errorf("duplicate primary key: %s", pk) } if err != ErrKeyNotFound { return 0, err } s.reads[key] = 0 } s.stagePut(table, key, nr) return rowid, nil } // UpdateByPK stages or durably applies a single-row update. func (s *Session) UpdateByPK(table, pk string, updateFn func(Row) (Row, error)) (Row, bool, error) { s.mu.Lock() if !s.inTx { s.mu.Unlock() return s.table.UpdateByPK(table, pk, updateFn) } defer s.mu.Unlock() schema, err := s.schema.GetSchema(table) if err != nil { return nil, false, err } key := s.table.dataKey(table, pk) row, err := s.getByPKLocked(table, pk) if err == ErrKeyNotFound { return nil, false, nil } if err != nil { return nil, false, err } oldRow := cloneRow(row) updates, err := s.runUpdateFn(updateFn, row) if err != nil { return nil, false, err } for name, value := range updates { for _, column := range schema.Columns { if strings.EqualFold(name, column.Name) { row[column.Name] = value break } } } s.stagePut(table, key, row) return oldRow, true, nil } // runUpdateFn releases the session lock while invoking the update callback so the // callback can run nested reads (e.g. a scalar subquery in SET) through the same // session without deadlocking on s.mu. A connection executes one statement at a // time, so the staged overlay cannot change while the callback runs. The lock is // re-acquired before returning. func (s *Session) runUpdateFn(updateFn func(Row) (Row, error), row Row) (Row, error) { s.mu.Unlock() defer s.mu.Lock() return updateFn(row) } // DeleteByPK stages or durably applies a single-row delete. func (s *Session) DeleteByPK(table, pk string) (Row, bool, error) { s.mu.Lock() defer s.mu.Unlock() if !s.inTx { return s.table.DeleteByPK(table, pk) } key := s.table.dataKey(table, pk) row, err := s.getByPKLocked(table, pk) if err == ErrKeyNotFound { return nil, false, nil } if err != nil { return nil, false, err } s.stageDelete(table, key, row) return row, true, nil } // getByPKLocked reads a row observing the overlay (caller holds s.mu). func (s *Session) getByPKLocked(table, pk string) (Row, error) { key := s.table.dataKey(table, pk) tl := strings.ToLower(table) if e, ok := s.overlay[tl][key]; ok { if e.absent { return nil, ErrKeyNotFound } return cloneRow(e.row), nil } row, lsn, err := s.table.getByPKWithLSN(table, pk) if err != nil { if err == ErrKeyNotFound { s.reads[key] = 0 return nil, ErrKeyNotFound } return nil, err } s.reads[key] = lsn return row, nil } // UpdateFunc stages or durably applies a scan-based update. func (s *Session) UpdateFunc(table string, updateFn func(Row) (Row, error), filter func(Row) bool) (int, error) { s.mu.Lock() if !s.inTx { s.mu.Unlock() return s.table.UpdateFunc(table, updateFn, filter) } defer s.mu.Unlock() schema, err := s.schema.GetSchema(table) if err != nil { return 0, err } rows, err := s.selectLocked(table, filter) if err != nil { return 0, err } count := 0 for _, row := range rows { oldRow := cloneRow(row) oldKey := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey])) updates, err := s.runUpdateFn(updateFn, row) if err != nil { return count, err } for name, value := range updates { for _, column := range schema.Columns { if strings.EqualFold(name, column.Name) { row[column.Name] = value break } } } newKey := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey])) if newKey != oldKey { // The primary key changed: move the staged row instead of leaving a // stale copy under the old key, and refuse to overwrite another row. if _, gerr := s.getByPKLocked(table, fmt.Sprintf("%v", row[schema.PrimaryKey])); gerr == nil { return count, fmt.Errorf("duplicate primary key: %v", row[schema.PrimaryKey]) } else if gerr != ErrKeyNotFound { return count, gerr } s.stageDelete(table, oldKey, oldRow) } s.stagePut(table, newKey, row) count++ } return count, nil } // Delete stages or durably applies a scan-based delete. func (s *Session) Delete(table string, filter func(Row) bool) (int, error) { s.mu.Lock() if !s.inTx { s.mu.Unlock() return s.table.Delete(table, filter) } defer s.mu.Unlock() schema, err := s.schema.GetSchema(table) if err != nil { return 0, err } rows, err := s.selectLocked(table, filter) if err != nil { return 0, err } count := 0 for _, row := range rows { key := s.table.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey])) s.stageDelete(table, key, row) count++ } return count, nil } // ClearIndex passes through to the durable table manager. func (s *Session) ClearIndex(indexName, tableName string, columns []string) error { return s.table.ClearIndex(indexName, tableName, columns) } // InvalidateCache passes through to the durable table manager. func (s *Session) InvalidateCache(table string) { s.table.InvalidateCache(table) } // BuildIndex passes through to the durable table manager. func (s *Session) BuildIndex(indexName, tableName string, columns []string) error { return s.table.BuildIndex(indexName, tableName, columns) }