You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

1407 lines
39 KiB

  1. /*
  2. Copyright 2015 Google LLC
  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. /*
  14. Package bttest contains test helpers for working with the bigtable package.
  15. To use a Server, create it, and then connect to it with no security:
  16. (The project/instance values are ignored.)
  17. srv, err := bttest.NewServer("localhost:0")
  18. ...
  19. conn, err := grpc.Dial(srv.Addr, grpc.WithInsecure())
  20. ...
  21. client, err := bigtable.NewClient(ctx, proj, instance,
  22. option.WithGRPCConn(conn))
  23. ...
  24. */
  25. package bttest // import "cloud.google.com/go/bigtable/bttest"
  26. import (
  27. "bytes"
  28. "context"
  29. "encoding/binary"
  30. "fmt"
  31. "log"
  32. "math/rand"
  33. "net"
  34. "regexp"
  35. "sort"
  36. "strings"
  37. "sync"
  38. "time"
  39. emptypb "github.com/golang/protobuf/ptypes/empty"
  40. "github.com/golang/protobuf/ptypes/wrappers"
  41. "github.com/google/btree"
  42. btapb "google.golang.org/genproto/googleapis/bigtable/admin/v2"
  43. btpb "google.golang.org/genproto/googleapis/bigtable/v2"
  44. "google.golang.org/genproto/googleapis/longrunning"
  45. statpb "google.golang.org/genproto/googleapis/rpc/status"
  46. "google.golang.org/grpc"
  47. "google.golang.org/grpc/codes"
  48. "google.golang.org/grpc/status"
  49. )
  50. const (
  51. // MilliSeconds field of the minimum valid Timestamp.
  52. minValidMilliSeconds = 0
  53. // MilliSeconds field of the max valid Timestamp.
  54. maxValidMilliSeconds = int64(time.Millisecond) * 253402300800
  55. )
  56. var (
  57. validLabelTransformer = regexp.MustCompile(`[a-z0-9\-]{1,15}`)
  58. )
  59. // Server is an in-memory Cloud Bigtable fake.
  60. // It is unauthenticated, and only a rough approximation.
  61. type Server struct {
  62. Addr string
  63. l net.Listener
  64. srv *grpc.Server
  65. s *server
  66. }
  67. // server is the real implementation of the fake.
  68. // It is a separate and unexported type so the API won't be cluttered with
  69. // methods that are only relevant to the fake's implementation.
  70. type server struct {
  71. mu sync.Mutex
  72. tables map[string]*table // keyed by fully qualified name
  73. gcc chan int // set when gcloop starts, closed when server shuts down
  74. // Any unimplemented methods will cause a panic.
  75. btapb.BigtableTableAdminServer
  76. btpb.BigtableServer
  77. }
  78. // NewServer creates a new Server.
  79. // The Server will be listening for gRPC connections, without TLS,
  80. // on the provided address. The resolved address is named by the Addr field.
  81. func NewServer(laddr string, opt ...grpc.ServerOption) (*Server, error) {
  82. l, err := net.Listen("tcp", laddr)
  83. if err != nil {
  84. return nil, err
  85. }
  86. s := &Server{
  87. Addr: l.Addr().String(),
  88. l: l,
  89. srv: grpc.NewServer(opt...),
  90. s: &server{
  91. tables: make(map[string]*table),
  92. },
  93. }
  94. btapb.RegisterBigtableTableAdminServer(s.srv, s.s)
  95. btpb.RegisterBigtableServer(s.srv, s.s)
  96. go s.srv.Serve(s.l)
  97. return s, nil
  98. }
  99. // Close shuts down the server.
  100. func (s *Server) Close() {
  101. s.s.mu.Lock()
  102. if s.s.gcc != nil {
  103. close(s.s.gcc)
  104. }
  105. s.s.mu.Unlock()
  106. s.srv.Stop()
  107. s.l.Close()
  108. }
  109. func (s *server) CreateTable(ctx context.Context, req *btapb.CreateTableRequest) (*btapb.Table, error) {
  110. tbl := req.Parent + "/tables/" + req.TableId
  111. s.mu.Lock()
  112. if _, ok := s.tables[tbl]; ok {
  113. s.mu.Unlock()
  114. return nil, status.Errorf(codes.AlreadyExists, "table %q already exists", tbl)
  115. }
  116. s.tables[tbl] = newTable(req)
  117. s.mu.Unlock()
  118. return &btapb.Table{Name: tbl}, nil
  119. }
  120. func (s *server) CreateTableFromSnapshot(context.Context, *btapb.CreateTableFromSnapshotRequest) (*longrunning.Operation, error) {
  121. return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
  122. }
  123. func (s *server) ListTables(ctx context.Context, req *btapb.ListTablesRequest) (*btapb.ListTablesResponse, error) {
  124. res := &btapb.ListTablesResponse{}
  125. prefix := req.Parent + "/tables/"
  126. s.mu.Lock()
  127. for tbl := range s.tables {
  128. if strings.HasPrefix(tbl, prefix) {
  129. res.Tables = append(res.Tables, &btapb.Table{Name: tbl})
  130. }
  131. }
  132. s.mu.Unlock()
  133. return res, nil
  134. }
  135. func (s *server) GetTable(ctx context.Context, req *btapb.GetTableRequest) (*btapb.Table, error) {
  136. tbl := req.Name
  137. s.mu.Lock()
  138. tblIns, ok := s.tables[tbl]
  139. s.mu.Unlock()
  140. if !ok {
  141. return nil, status.Errorf(codes.NotFound, "table %q not found", tbl)
  142. }
  143. return &btapb.Table{
  144. Name: tbl,
  145. ColumnFamilies: toColumnFamilies(tblIns.columnFamilies()),
  146. }, nil
  147. }
  148. func (s *server) DeleteTable(ctx context.Context, req *btapb.DeleteTableRequest) (*emptypb.Empty, error) {
  149. s.mu.Lock()
  150. defer s.mu.Unlock()
  151. if _, ok := s.tables[req.Name]; !ok {
  152. return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
  153. }
  154. delete(s.tables, req.Name)
  155. return &emptypb.Empty{}, nil
  156. }
  157. func (s *server) ModifyColumnFamilies(ctx context.Context, req *btapb.ModifyColumnFamiliesRequest) (*btapb.Table, error) {
  158. s.mu.Lock()
  159. tbl, ok := s.tables[req.Name]
  160. s.mu.Unlock()
  161. if !ok {
  162. return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
  163. }
  164. tbl.mu.Lock()
  165. defer tbl.mu.Unlock()
  166. for _, mod := range req.Modifications {
  167. if create := mod.GetCreate(); create != nil {
  168. if _, ok := tbl.families[mod.Id]; ok {
  169. return nil, status.Errorf(codes.AlreadyExists, "family %q already exists", mod.Id)
  170. }
  171. newcf := &columnFamily{
  172. name: req.Name + "/columnFamilies/" + mod.Id,
  173. order: tbl.counter,
  174. gcRule: create.GcRule,
  175. }
  176. tbl.counter++
  177. tbl.families[mod.Id] = newcf
  178. } else if mod.GetDrop() {
  179. if _, ok := tbl.families[mod.Id]; !ok {
  180. return nil, fmt.Errorf("can't delete unknown family %q", mod.Id)
  181. }
  182. delete(tbl.families, mod.Id)
  183. } else if modify := mod.GetUpdate(); modify != nil {
  184. if _, ok := tbl.families[mod.Id]; !ok {
  185. return nil, fmt.Errorf("no such family %q", mod.Id)
  186. }
  187. newcf := &columnFamily{
  188. name: req.Name + "/columnFamilies/" + mod.Id,
  189. gcRule: modify.GcRule,
  190. }
  191. // assume that we ALWAYS want to replace by the new setting
  192. // we may need partial update through
  193. tbl.families[mod.Id] = newcf
  194. }
  195. }
  196. s.needGC()
  197. return &btapb.Table{
  198. Name: req.Name,
  199. ColumnFamilies: toColumnFamilies(tbl.families),
  200. Granularity: btapb.Table_TimestampGranularity(btapb.Table_MILLIS),
  201. }, nil
  202. }
  203. func (s *server) DropRowRange(ctx context.Context, req *btapb.DropRowRangeRequest) (*emptypb.Empty, error) {
  204. s.mu.Lock()
  205. defer s.mu.Unlock()
  206. tbl, ok := s.tables[req.Name]
  207. if !ok {
  208. return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
  209. }
  210. if req.GetDeleteAllDataFromTable() {
  211. tbl.rows = btree.New(btreeDegree)
  212. } else {
  213. // Delete rows by prefix.
  214. prefixBytes := req.GetRowKeyPrefix()
  215. if prefixBytes == nil {
  216. return nil, fmt.Errorf("missing row key prefix")
  217. }
  218. prefix := string(prefixBytes)
  219. // The BTree does not specify what happens if rows are deleted during
  220. // iteration, and it provides no "delete range" method.
  221. // So we collect the rows first, then delete them one by one.
  222. var rowsToDelete []*row
  223. tbl.rows.AscendGreaterOrEqual(btreeKey(prefix), func(i btree.Item) bool {
  224. r := i.(*row)
  225. if strings.HasPrefix(r.key, prefix) {
  226. rowsToDelete = append(rowsToDelete, r)
  227. return true
  228. }
  229. return false // stop iteration
  230. })
  231. for _, r := range rowsToDelete {
  232. tbl.rows.Delete(r)
  233. }
  234. }
  235. return &emptypb.Empty{}, nil
  236. }
  237. func (s *server) GenerateConsistencyToken(ctx context.Context, req *btapb.GenerateConsistencyTokenRequest) (*btapb.GenerateConsistencyTokenResponse, error) {
  238. // Check that the table exists.
  239. _, ok := s.tables[req.Name]
  240. if !ok {
  241. return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
  242. }
  243. return &btapb.GenerateConsistencyTokenResponse{
  244. ConsistencyToken: "TokenFor-" + req.Name,
  245. }, nil
  246. }
  247. func (s *server) CheckConsistency(ctx context.Context, req *btapb.CheckConsistencyRequest) (*btapb.CheckConsistencyResponse, error) {
  248. // Check that the table exists.
  249. _, ok := s.tables[req.Name]
  250. if !ok {
  251. return nil, status.Errorf(codes.NotFound, "table %q not found", req.Name)
  252. }
  253. // Check this is the right token.
  254. if req.ConsistencyToken != "TokenFor-"+req.Name {
  255. return nil, status.Errorf(codes.InvalidArgument, "token %q not valid", req.ConsistencyToken)
  256. }
  257. // Single cluster instances are always consistent.
  258. return &btapb.CheckConsistencyResponse{
  259. Consistent: true,
  260. }, nil
  261. }
  262. func (s *server) SnapshotTable(context.Context, *btapb.SnapshotTableRequest) (*longrunning.Operation, error) {
  263. return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
  264. }
  265. func (s *server) GetSnapshot(context.Context, *btapb.GetSnapshotRequest) (*btapb.Snapshot, error) {
  266. return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
  267. }
  268. func (s *server) ListSnapshots(context.Context, *btapb.ListSnapshotsRequest) (*btapb.ListSnapshotsResponse, error) {
  269. return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
  270. }
  271. func (s *server) DeleteSnapshot(context.Context, *btapb.DeleteSnapshotRequest) (*emptypb.Empty, error) {
  272. return nil, status.Errorf(codes.Unimplemented, "the emulator does not currently support snapshots")
  273. }
  274. func (s *server) ReadRows(req *btpb.ReadRowsRequest, stream btpb.Bigtable_ReadRowsServer) error {
  275. s.mu.Lock()
  276. tbl, ok := s.tables[req.TableName]
  277. s.mu.Unlock()
  278. if !ok {
  279. return status.Errorf(codes.NotFound, "table %q not found", req.TableName)
  280. }
  281. // Rows to read can be specified by a set of row keys and/or a set of row ranges.
  282. // Output is a stream of sorted, de-duped rows.
  283. tbl.mu.RLock()
  284. rowSet := make(map[string]*row)
  285. addRow := func(i btree.Item) bool {
  286. r := i.(*row)
  287. rowSet[r.key] = r
  288. return true
  289. }
  290. if req.Rows != nil &&
  291. len(req.Rows.RowKeys)+len(req.Rows.RowRanges) > 0 {
  292. // Add the explicitly given keys
  293. for _, key := range req.Rows.RowKeys {
  294. k := string(key)
  295. if i := tbl.rows.Get(btreeKey(k)); i != nil {
  296. addRow(i)
  297. }
  298. }
  299. // Add keys from row ranges
  300. for _, rr := range req.Rows.RowRanges {
  301. var start, end string
  302. switch sk := rr.StartKey.(type) {
  303. case *btpb.RowRange_StartKeyClosed:
  304. start = string(sk.StartKeyClosed)
  305. case *btpb.RowRange_StartKeyOpen:
  306. start = string(sk.StartKeyOpen) + "\x00"
  307. }
  308. switch ek := rr.EndKey.(type) {
  309. case *btpb.RowRange_EndKeyClosed:
  310. end = string(ek.EndKeyClosed) + "\x00"
  311. case *btpb.RowRange_EndKeyOpen:
  312. end = string(ek.EndKeyOpen)
  313. }
  314. switch {
  315. case start == "" && end == "":
  316. tbl.rows.Ascend(addRow) // all rows
  317. case start == "":
  318. tbl.rows.AscendLessThan(btreeKey(end), addRow)
  319. case end == "":
  320. tbl.rows.AscendGreaterOrEqual(btreeKey(start), addRow)
  321. default:
  322. tbl.rows.AscendRange(btreeKey(start), btreeKey(end), addRow)
  323. }
  324. }
  325. } else {
  326. // Read all rows
  327. tbl.rows.Ascend(addRow)
  328. }
  329. tbl.mu.RUnlock()
  330. rows := make([]*row, 0, len(rowSet))
  331. for _, r := range rowSet {
  332. r.mu.Lock()
  333. fams := len(r.families)
  334. r.mu.Unlock()
  335. if fams != 0 {
  336. rows = append(rows, r)
  337. }
  338. }
  339. sort.Sort(byRowKey(rows))
  340. limit := int(req.RowsLimit)
  341. count := 0
  342. for _, r := range rows {
  343. if limit > 0 && count >= limit {
  344. return nil
  345. }
  346. streamed, err := streamRow(stream, r, req.Filter)
  347. if err != nil {
  348. return err
  349. }
  350. if streamed {
  351. count++
  352. }
  353. }
  354. return nil
  355. }
  356. // streamRow filters the given row and sends it via the given stream.
  357. // Returns true if at least one cell matched the filter and was streamed, false otherwise.
  358. func streamRow(stream btpb.Bigtable_ReadRowsServer, r *row, f *btpb.RowFilter) (bool, error) {
  359. r.mu.Lock()
  360. nr := r.copy()
  361. r.mu.Unlock()
  362. r = nr
  363. match, err := filterRow(f, r)
  364. if err != nil {
  365. return false, err
  366. }
  367. if !match {
  368. return false, nil
  369. }
  370. rrr := &btpb.ReadRowsResponse{}
  371. families := r.sortedFamilies()
  372. for _, fam := range families {
  373. for _, colName := range fam.colNames {
  374. cells := fam.cells[colName]
  375. if len(cells) == 0 {
  376. continue
  377. }
  378. for _, cell := range cells {
  379. rrr.Chunks = append(rrr.Chunks, &btpb.ReadRowsResponse_CellChunk{
  380. RowKey: []byte(r.key),
  381. FamilyName: &wrappers.StringValue{Value: fam.name},
  382. Qualifier: &wrappers.BytesValue{Value: []byte(colName)},
  383. TimestampMicros: cell.ts,
  384. Value: cell.value,
  385. Labels: cell.labels,
  386. })
  387. }
  388. }
  389. }
  390. // We can't have a cell with just COMMIT set, which would imply a new empty cell.
  391. // So modify the last cell to have the COMMIT flag set.
  392. if len(rrr.Chunks) > 0 {
  393. rrr.Chunks[len(rrr.Chunks)-1].RowStatus = &btpb.ReadRowsResponse_CellChunk_CommitRow{CommitRow: true}
  394. }
  395. return true, stream.Send(rrr)
  396. }
  397. // filterRow modifies a row with the given filter. Returns true if at least one cell from the row matches,
  398. // false otherwise. If a filter is invalid, filterRow returns false and an error.
  399. func filterRow(f *btpb.RowFilter, r *row) (bool, error) {
  400. if f == nil {
  401. return true, nil
  402. }
  403. // Handle filters that apply beyond just including/excluding cells.
  404. switch f := f.Filter.(type) {
  405. case *btpb.RowFilter_BlockAllFilter:
  406. return !f.BlockAllFilter, nil
  407. case *btpb.RowFilter_PassAllFilter:
  408. return f.PassAllFilter, nil
  409. case *btpb.RowFilter_Chain_:
  410. for _, sub := range f.Chain.Filters {
  411. match, err := filterRow(sub, r)
  412. if err != nil {
  413. return false, err
  414. }
  415. if !match {
  416. return false, nil
  417. }
  418. }
  419. return true, nil
  420. case *btpb.RowFilter_Interleave_:
  421. srs := make([]*row, 0, len(f.Interleave.Filters))
  422. for _, sub := range f.Interleave.Filters {
  423. sr := r.copy()
  424. filterRow(sub, sr)
  425. srs = append(srs, sr)
  426. }
  427. // merge
  428. // TODO(dsymonds): is this correct?
  429. r.families = make(map[string]*family)
  430. for _, sr := range srs {
  431. for _, fam := range sr.families {
  432. f := r.getOrCreateFamily(fam.name, fam.order)
  433. for colName, cs := range fam.cells {
  434. f.cells[colName] = append(f.cellsByColumn(colName), cs...)
  435. }
  436. }
  437. }
  438. var count int
  439. for _, fam := range r.families {
  440. for _, cs := range fam.cells {
  441. sort.Sort(byDescTS(cs))
  442. count += len(cs)
  443. }
  444. }
  445. return count > 0, nil
  446. case *btpb.RowFilter_CellsPerColumnLimitFilter:
  447. lim := int(f.CellsPerColumnLimitFilter)
  448. for _, fam := range r.families {
  449. for col, cs := range fam.cells {
  450. if len(cs) > lim {
  451. fam.cells[col] = cs[:lim]
  452. }
  453. }
  454. }
  455. return true, nil
  456. case *btpb.RowFilter_Condition_:
  457. match, err := filterRow(f.Condition.PredicateFilter, r.copy())
  458. if err != nil {
  459. return false, err
  460. }
  461. if match {
  462. if f.Condition.TrueFilter == nil {
  463. return false, nil
  464. }
  465. return filterRow(f.Condition.TrueFilter, r)
  466. }
  467. if f.Condition.FalseFilter == nil {
  468. return false, nil
  469. }
  470. return filterRow(f.Condition.FalseFilter, r)
  471. case *btpb.RowFilter_RowKeyRegexFilter:
  472. rx, err := newRegexp(f.RowKeyRegexFilter)
  473. if err != nil {
  474. return false, status.Errorf(codes.InvalidArgument, "Error in field 'rowkey_regex_filter' : %v", err)
  475. }
  476. if !rx.MatchString(r.key) {
  477. return false, nil
  478. }
  479. case *btpb.RowFilter_CellsPerRowLimitFilter:
  480. // Grab the first n cells in the row.
  481. lim := int(f.CellsPerRowLimitFilter)
  482. for _, fam := range r.families {
  483. for _, col := range fam.colNames {
  484. cs := fam.cells[col]
  485. if len(cs) > lim {
  486. fam.cells[col] = cs[:lim]
  487. lim = 0
  488. } else {
  489. lim -= len(cs)
  490. }
  491. }
  492. }
  493. return true, nil
  494. case *btpb.RowFilter_CellsPerRowOffsetFilter:
  495. // Skip the first n cells in the row.
  496. offset := int(f.CellsPerRowOffsetFilter)
  497. for _, fam := range r.families {
  498. for _, col := range fam.colNames {
  499. cs := fam.cells[col]
  500. if len(cs) > offset {
  501. fam.cells[col] = cs[offset:]
  502. offset = 0
  503. return true, nil
  504. }
  505. fam.cells[col] = cs[:0]
  506. offset -= len(cs)
  507. }
  508. }
  509. return true, nil
  510. case *btpb.RowFilter_RowSampleFilter:
  511. // The row sample filter "matches all cells from a row with probability
  512. // p, and matches no cells from the row with probability 1-p."
  513. // See https://github.com/googleapis/googleapis/blob/master/google/bigtable/v2/data.proto
  514. if f.RowSampleFilter <= 0.0 || f.RowSampleFilter >= 1.0 {
  515. return false, status.Error(codes.InvalidArgument, "row_sample_filter argument must be between 0.0 and 1.0")
  516. }
  517. return randFloat() < f.RowSampleFilter, nil
  518. }
  519. // Any other case, operate on a per-cell basis.
  520. cellCount := 0
  521. for _, fam := range r.families {
  522. for colName, cs := range fam.cells {
  523. filtered, err := filterCells(f, fam.name, colName, cs)
  524. if err != nil {
  525. return false, err
  526. }
  527. fam.cells[colName] = filtered
  528. cellCount += len(fam.cells[colName])
  529. }
  530. }
  531. return cellCount > 0, nil
  532. }
  533. var randFloat = rand.Float64
  534. func filterCells(f *btpb.RowFilter, fam, col string, cs []cell) ([]cell, error) {
  535. var ret []cell
  536. for _, cell := range cs {
  537. include, err := includeCell(f, fam, col, cell)
  538. if err != nil {
  539. return nil, err
  540. }
  541. if include {
  542. cell, err = modifyCell(f, cell)
  543. if err != nil {
  544. return nil, err
  545. }
  546. ret = append(ret, cell)
  547. }
  548. }
  549. return ret, nil
  550. }
  551. func modifyCell(f *btpb.RowFilter, c cell) (cell, error) {
  552. if f == nil {
  553. return c, nil
  554. }
  555. // Consider filters that may modify the cell contents
  556. switch filter := f.Filter.(type) {
  557. case *btpb.RowFilter_StripValueTransformer:
  558. return cell{ts: c.ts}, nil
  559. case *btpb.RowFilter_ApplyLabelTransformer:
  560. if !validLabelTransformer.MatchString(filter.ApplyLabelTransformer) {
  561. return cell{}, status.Errorf(
  562. codes.InvalidArgument,
  563. `apply_label_transformer must match RE2([a-z0-9\-]+), but found %v`,
  564. filter.ApplyLabelTransformer,
  565. )
  566. }
  567. return cell{ts: c.ts, value: c.value, labels: []string{filter.ApplyLabelTransformer}}, nil
  568. default:
  569. return c, nil
  570. }
  571. }
  572. func includeCell(f *btpb.RowFilter, fam, col string, cell cell) (bool, error) {
  573. if f == nil {
  574. return true, nil
  575. }
  576. // TODO(dsymonds): Implement many more filters.
  577. switch f := f.Filter.(type) {
  578. case *btpb.RowFilter_CellsPerColumnLimitFilter:
  579. // Don't log, row-level filter
  580. return true, nil
  581. case *btpb.RowFilter_RowKeyRegexFilter:
  582. // Don't log, row-level filter
  583. return true, nil
  584. case *btpb.RowFilter_StripValueTransformer:
  585. // Don't log, cell-modifying filter
  586. return true, nil
  587. case *btpb.RowFilter_ApplyLabelTransformer:
  588. // Don't log, cell-modifying filter
  589. return true, nil
  590. default:
  591. log.Printf("WARNING: don't know how to handle filter of type %T (ignoring it)", f)
  592. return true, nil
  593. case *btpb.RowFilter_FamilyNameRegexFilter:
  594. rx, err := newRegexp([]byte(f.FamilyNameRegexFilter))
  595. if err != nil {
  596. return false, status.Errorf(codes.InvalidArgument, "Error in field 'family_name_regex_filter' : %v", err)
  597. }
  598. return rx.MatchString(fam), nil
  599. case *btpb.RowFilter_ColumnQualifierRegexFilter:
  600. rx, err := newRegexp(f.ColumnQualifierRegexFilter)
  601. if err != nil {
  602. return false, status.Errorf(codes.InvalidArgument, "Error in field 'column_qualifier_regex_filter' : %v", err)
  603. }
  604. return rx.MatchString(toUTF8([]byte(col))), nil
  605. case *btpb.RowFilter_ValueRegexFilter:
  606. rx, err := newRegexp(f.ValueRegexFilter)
  607. if err != nil {
  608. return false, status.Errorf(codes.InvalidArgument, "Error in field 'value_regex_filter' : %v", err)
  609. }
  610. return rx.Match(cell.value), nil
  611. case *btpb.RowFilter_ColumnRangeFilter:
  612. if fam != f.ColumnRangeFilter.FamilyName {
  613. return false, nil
  614. }
  615. // Start qualifier defaults to empty string closed
  616. inRangeStart := func() bool { return col >= "" }
  617. switch sq := f.ColumnRangeFilter.StartQualifier.(type) {
  618. case *btpb.ColumnRange_StartQualifierOpen:
  619. inRangeStart = func() bool { return col > string(sq.StartQualifierOpen) }
  620. case *btpb.ColumnRange_StartQualifierClosed:
  621. inRangeStart = func() bool { return col >= string(sq.StartQualifierClosed) }
  622. }
  623. // End qualifier defaults to no upper boundary
  624. inRangeEnd := func() bool { return true }
  625. switch eq := f.ColumnRangeFilter.EndQualifier.(type) {
  626. case *btpb.ColumnRange_EndQualifierClosed:
  627. inRangeEnd = func() bool { return col <= string(eq.EndQualifierClosed) }
  628. case *btpb.ColumnRange_EndQualifierOpen:
  629. inRangeEnd = func() bool { return col < string(eq.EndQualifierOpen) }
  630. }
  631. return inRangeStart() && inRangeEnd(), nil
  632. case *btpb.RowFilter_TimestampRangeFilter:
  633. // Lower bound is inclusive and defaults to 0, upper bound is exclusive and defaults to infinity.
  634. return cell.ts >= f.TimestampRangeFilter.StartTimestampMicros &&
  635. (f.TimestampRangeFilter.EndTimestampMicros == 0 || cell.ts < f.TimestampRangeFilter.EndTimestampMicros), nil
  636. case *btpb.RowFilter_ValueRangeFilter:
  637. v := cell.value
  638. // Start value defaults to empty string closed
  639. inRangeStart := func() bool { return bytes.Compare(v, []byte{}) >= 0 }
  640. switch sv := f.ValueRangeFilter.StartValue.(type) {
  641. case *btpb.ValueRange_StartValueOpen:
  642. inRangeStart = func() bool { return bytes.Compare(v, sv.StartValueOpen) > 0 }
  643. case *btpb.ValueRange_StartValueClosed:
  644. inRangeStart = func() bool { return bytes.Compare(v, sv.StartValueClosed) >= 0 }
  645. }
  646. // End value defaults to no upper boundary
  647. inRangeEnd := func() bool { return true }
  648. switch ev := f.ValueRangeFilter.EndValue.(type) {
  649. case *btpb.ValueRange_EndValueClosed:
  650. inRangeEnd = func() bool { return bytes.Compare(v, ev.EndValueClosed) <= 0 }
  651. case *btpb.ValueRange_EndValueOpen:
  652. inRangeEnd = func() bool { return bytes.Compare(v, ev.EndValueOpen) < 0 }
  653. }
  654. return inRangeStart() && inRangeEnd(), nil
  655. }
  656. }
  657. func toUTF8(bs []byte) string {
  658. var rs []rune
  659. for _, b := range bs {
  660. rs = append(rs, rune(b))
  661. }
  662. return string(rs)
  663. }
  664. func newRegexp(patBytes []byte) (*regexp.Regexp, error) {
  665. pat := toUTF8(patBytes)
  666. re, err := regexp.Compile("^" + pat + "$") // match entire target
  667. if err != nil {
  668. log.Printf("Bad pattern %q: %v", pat, err)
  669. }
  670. return re, err
  671. }
  672. func (s *server) MutateRow(ctx context.Context, req *btpb.MutateRowRequest) (*btpb.MutateRowResponse, error) {
  673. s.mu.Lock()
  674. tbl, ok := s.tables[req.TableName]
  675. s.mu.Unlock()
  676. if !ok {
  677. return nil, status.Errorf(codes.NotFound, "table %q not found", req.TableName)
  678. }
  679. fs := tbl.columnFamilies()
  680. r := tbl.mutableRow(string(req.RowKey))
  681. r.mu.Lock()
  682. defer r.mu.Unlock()
  683. if err := applyMutations(tbl, r, req.Mutations, fs); err != nil {
  684. return nil, err
  685. }
  686. return &btpb.MutateRowResponse{}, nil
  687. }
  688. func (s *server) MutateRows(req *btpb.MutateRowsRequest, stream btpb.Bigtable_MutateRowsServer) error {
  689. s.mu.Lock()
  690. tbl, ok := s.tables[req.TableName]
  691. s.mu.Unlock()
  692. if !ok {
  693. return status.Errorf(codes.NotFound, "table %q not found", req.TableName)
  694. }
  695. res := &btpb.MutateRowsResponse{Entries: make([]*btpb.MutateRowsResponse_Entry, len(req.Entries))}
  696. fs := tbl.columnFamilies()
  697. for i, entry := range req.Entries {
  698. r := tbl.mutableRow(string(entry.RowKey))
  699. r.mu.Lock()
  700. code, msg := int32(codes.OK), ""
  701. if err := applyMutations(tbl, r, entry.Mutations, fs); err != nil {
  702. code = int32(codes.Internal)
  703. msg = err.Error()
  704. }
  705. res.Entries[i] = &btpb.MutateRowsResponse_Entry{
  706. Index: int64(i),
  707. Status: &statpb.Status{Code: code, Message: msg},
  708. }
  709. r.mu.Unlock()
  710. }
  711. return stream.Send(res)
  712. }
  713. func (s *server) CheckAndMutateRow(ctx context.Context, req *btpb.CheckAndMutateRowRequest) (*btpb.CheckAndMutateRowResponse, error) {
  714. s.mu.Lock()
  715. tbl, ok := s.tables[req.TableName]
  716. s.mu.Unlock()
  717. if !ok {
  718. return nil, status.Errorf(codes.NotFound, "table %q not found", req.TableName)
  719. }
  720. res := &btpb.CheckAndMutateRowResponse{}
  721. fs := tbl.columnFamilies()
  722. r := tbl.mutableRow(string(req.RowKey))
  723. r.mu.Lock()
  724. defer r.mu.Unlock()
  725. // Figure out which mutation to apply.
  726. whichMut := false
  727. if req.PredicateFilter == nil {
  728. // Use true_mutations iff row contains any cells.
  729. whichMut = !r.isEmpty()
  730. } else {
  731. // Use true_mutations iff any cells in the row match the filter.
  732. // TODO(dsymonds): This could be cheaper.
  733. nr := r.copy()
  734. filterRow(req.PredicateFilter, nr)
  735. whichMut = !nr.isEmpty()
  736. }
  737. res.PredicateMatched = whichMut
  738. muts := req.FalseMutations
  739. if whichMut {
  740. muts = req.TrueMutations
  741. }
  742. if err := applyMutations(tbl, r, muts, fs); err != nil {
  743. return nil, err
  744. }
  745. return res, nil
  746. }
  747. // applyMutations applies a sequence of mutations to a row.
  748. // fam should be a snapshot of the keys of tbl.families.
  749. // It assumes r.mu is locked.
  750. func applyMutations(tbl *table, r *row, muts []*btpb.Mutation, fs map[string]*columnFamily) error {
  751. for _, mut := range muts {
  752. switch mut := mut.Mutation.(type) {
  753. default:
  754. return fmt.Errorf("can't handle mutation type %T", mut)
  755. case *btpb.Mutation_SetCell_:
  756. set := mut.SetCell
  757. if _, ok := fs[set.FamilyName]; !ok {
  758. return fmt.Errorf("unknown family %q", set.FamilyName)
  759. }
  760. ts := set.TimestampMicros
  761. if ts == -1 { // bigtable.ServerTime
  762. ts = newTimestamp()
  763. }
  764. if !tbl.validTimestamp(ts) {
  765. return fmt.Errorf("invalid timestamp %d", ts)
  766. }
  767. fam := set.FamilyName
  768. col := string(set.ColumnQualifier)
  769. newCell := cell{ts: ts, value: set.Value}
  770. f := r.getOrCreateFamily(fam, fs[fam].order)
  771. f.cells[col] = appendOrReplaceCell(f.cellsByColumn(col), newCell)
  772. case *btpb.Mutation_DeleteFromColumn_:
  773. del := mut.DeleteFromColumn
  774. if _, ok := fs[del.FamilyName]; !ok {
  775. return fmt.Errorf("unknown family %q", del.FamilyName)
  776. }
  777. fam := del.FamilyName
  778. col := string(del.ColumnQualifier)
  779. if _, ok := r.families[fam]; ok {
  780. cs := r.families[fam].cells[col]
  781. if del.TimeRange != nil {
  782. tsr := del.TimeRange
  783. if !tbl.validTimestamp(tsr.StartTimestampMicros) {
  784. return fmt.Errorf("invalid timestamp %d", tsr.StartTimestampMicros)
  785. }
  786. if !tbl.validTimestamp(tsr.EndTimestampMicros) && tsr.EndTimestampMicros != 0 {
  787. return fmt.Errorf("invalid timestamp %d", tsr.EndTimestampMicros)
  788. }
  789. if tsr.StartTimestampMicros >= tsr.EndTimestampMicros && tsr.EndTimestampMicros != 0 {
  790. return fmt.Errorf("inverted or invalid timestamp range [%d, %d]", tsr.StartTimestampMicros, tsr.EndTimestampMicros)
  791. }
  792. // Find half-open interval to remove.
  793. // Cells are in descending timestamp order,
  794. // so the predicates to sort.Search are inverted.
  795. si, ei := 0, len(cs)
  796. if tsr.StartTimestampMicros > 0 {
  797. ei = sort.Search(len(cs), func(i int) bool { return cs[i].ts < tsr.StartTimestampMicros })
  798. }
  799. if tsr.EndTimestampMicros > 0 {
  800. si = sort.Search(len(cs), func(i int) bool { return cs[i].ts < tsr.EndTimestampMicros })
  801. }
  802. if si < ei {
  803. copy(cs[si:], cs[ei:])
  804. cs = cs[:len(cs)-(ei-si)]
  805. }
  806. } else {
  807. cs = nil
  808. }
  809. if len(cs) == 0 {
  810. delete(r.families[fam].cells, col)
  811. colNames := r.families[fam].colNames
  812. i := sort.Search(len(colNames), func(i int) bool { return colNames[i] >= col })
  813. if i < len(colNames) && colNames[i] == col {
  814. r.families[fam].colNames = append(colNames[:i], colNames[i+1:]...)
  815. }
  816. if len(r.families[fam].cells) == 0 {
  817. delete(r.families, fam)
  818. }
  819. } else {
  820. r.families[fam].cells[col] = cs
  821. }
  822. }
  823. case *btpb.Mutation_DeleteFromRow_:
  824. r.families = make(map[string]*family)
  825. case *btpb.Mutation_DeleteFromFamily_:
  826. fampre := mut.DeleteFromFamily.FamilyName
  827. delete(r.families, fampre)
  828. }
  829. }
  830. return nil
  831. }
  832. func maxTimestamp(x, y int64) int64 {
  833. if x > y {
  834. return x
  835. }
  836. return y
  837. }
  838. func newTimestamp() int64 {
  839. ts := time.Now().UnixNano() / 1e3
  840. ts -= ts % 1000 // round to millisecond granularity
  841. return ts
  842. }
  843. func appendOrReplaceCell(cs []cell, newCell cell) []cell {
  844. replaced := false
  845. for i, cell := range cs {
  846. if cell.ts == newCell.ts {
  847. cs[i] = newCell
  848. replaced = true
  849. break
  850. }
  851. }
  852. if !replaced {
  853. cs = append(cs, newCell)
  854. }
  855. sort.Sort(byDescTS(cs))
  856. return cs
  857. }
  858. func (s *server) ReadModifyWriteRow(ctx context.Context, req *btpb.ReadModifyWriteRowRequest) (*btpb.ReadModifyWriteRowResponse, error) {
  859. s.mu.Lock()
  860. tbl, ok := s.tables[req.TableName]
  861. s.mu.Unlock()
  862. if !ok {
  863. return nil, status.Errorf(codes.NotFound, "table %q not found", req.TableName)
  864. }
  865. fs := tbl.columnFamilies()
  866. rowKey := string(req.RowKey)
  867. r := tbl.mutableRow(rowKey)
  868. resultRow := newRow(rowKey) // copy of updated cells
  869. // This must be done before the row lock, acquired below, is released.
  870. r.mu.Lock()
  871. defer r.mu.Unlock()
  872. // Assume all mutations apply to the most recent version of the cell.
  873. // TODO(dsymonds): Verify this assumption and document it in the proto.
  874. for _, rule := range req.Rules {
  875. if _, ok := fs[rule.FamilyName]; !ok {
  876. return nil, fmt.Errorf("unknown family %q", rule.FamilyName)
  877. }
  878. fam := rule.FamilyName
  879. col := string(rule.ColumnQualifier)
  880. isEmpty := false
  881. f := r.getOrCreateFamily(fam, fs[fam].order)
  882. cs := f.cells[col]
  883. isEmpty = len(cs) == 0
  884. ts := newTimestamp()
  885. var newCell, prevCell cell
  886. if !isEmpty {
  887. cells := r.families[fam].cells[col]
  888. prevCell = cells[0]
  889. // ts is the max of now or the prev cell's timestamp in case the
  890. // prev cell is in the future
  891. ts = maxTimestamp(ts, prevCell.ts)
  892. }
  893. switch rule := rule.Rule.(type) {
  894. default:
  895. return nil, fmt.Errorf("unknown RMW rule oneof %T", rule)
  896. case *btpb.ReadModifyWriteRule_AppendValue:
  897. newCell = cell{ts: ts, value: append(prevCell.value, rule.AppendValue...)}
  898. case *btpb.ReadModifyWriteRule_IncrementAmount:
  899. var v int64
  900. if !isEmpty {
  901. prevVal := prevCell.value
  902. if len(prevVal) != 8 {
  903. return nil, fmt.Errorf("increment on non-64-bit value")
  904. }
  905. v = int64(binary.BigEndian.Uint64(prevVal))
  906. }
  907. v += rule.IncrementAmount
  908. var val [8]byte
  909. binary.BigEndian.PutUint64(val[:], uint64(v))
  910. newCell = cell{ts: ts, value: val[:]}
  911. }
  912. // Store the new cell
  913. f.cells[col] = appendOrReplaceCell(f.cellsByColumn(col), newCell)
  914. // Store a copy for the result row
  915. resultFamily := resultRow.getOrCreateFamily(fam, fs[fam].order)
  916. resultFamily.cellsByColumn(col) // create the column
  917. resultFamily.cells[col] = []cell{newCell} // overwrite the cells
  918. }
  919. // Build the response using the result row
  920. res := &btpb.Row{
  921. Key: req.RowKey,
  922. Families: make([]*btpb.Family, len(resultRow.families)),
  923. }
  924. for i, family := range resultRow.sortedFamilies() {
  925. res.Families[i] = &btpb.Family{
  926. Name: family.name,
  927. Columns: make([]*btpb.Column, len(family.colNames)),
  928. }
  929. for j, colName := range family.colNames {
  930. res.Families[i].Columns[j] = &btpb.Column{
  931. Qualifier: []byte(colName),
  932. Cells: []*btpb.Cell{{
  933. TimestampMicros: family.cells[colName][0].ts,
  934. Value: family.cells[colName][0].value,
  935. }},
  936. }
  937. }
  938. }
  939. return &btpb.ReadModifyWriteRowResponse{Row: res}, nil
  940. }
  941. func (s *server) SampleRowKeys(req *btpb.SampleRowKeysRequest, stream btpb.Bigtable_SampleRowKeysServer) error {
  942. s.mu.Lock()
  943. tbl, ok := s.tables[req.TableName]
  944. s.mu.Unlock()
  945. if !ok {
  946. return status.Errorf(codes.NotFound, "table %q not found", req.TableName)
  947. }
  948. tbl.mu.RLock()
  949. defer tbl.mu.RUnlock()
  950. // The return value of SampleRowKeys is very loosely defined. Return at least the
  951. // final row key in the table and choose other row keys randomly.
  952. var offset int64
  953. var err error
  954. i := 0
  955. tbl.rows.Ascend(func(it btree.Item) bool {
  956. row := it.(*row)
  957. if i == tbl.rows.Len()-1 || rand.Int31n(100) == 0 {
  958. resp := &btpb.SampleRowKeysResponse{
  959. RowKey: []byte(row.key),
  960. OffsetBytes: offset,
  961. }
  962. err = stream.Send(resp)
  963. if err != nil {
  964. return false
  965. }
  966. }
  967. offset += int64(row.size())
  968. i++
  969. return true
  970. })
  971. return err
  972. }
  973. // needGC is invoked whenever the server needs gcloop running.
  974. func (s *server) needGC() {
  975. s.mu.Lock()
  976. if s.gcc == nil {
  977. s.gcc = make(chan int)
  978. go s.gcloop(s.gcc)
  979. }
  980. s.mu.Unlock()
  981. }
  982. func (s *server) gcloop(done <-chan int) {
  983. const (
  984. minWait = 500 // ms
  985. maxWait = 1500 // ms
  986. )
  987. for {
  988. // Wait for a random time interval.
  989. d := time.Duration(minWait+rand.Intn(maxWait-minWait)) * time.Millisecond
  990. select {
  991. case <-time.After(d):
  992. case <-done:
  993. return // server has been closed
  994. }
  995. // Do a GC pass over all tables.
  996. var tables []*table
  997. s.mu.Lock()
  998. for _, tbl := range s.tables {
  999. tables = append(tables, tbl)
  1000. }
  1001. s.mu.Unlock()
  1002. for _, tbl := range tables {
  1003. tbl.gc()
  1004. }
  1005. }
  1006. }
  1007. type table struct {
  1008. mu sync.RWMutex
  1009. counter uint64 // increment by 1 when a new family is created
  1010. families map[string]*columnFamily // keyed by plain family name
  1011. rows *btree.BTree // indexed by row key
  1012. }
  1013. const btreeDegree = 16
  1014. func newTable(ctr *btapb.CreateTableRequest) *table {
  1015. fams := make(map[string]*columnFamily)
  1016. c := uint64(0)
  1017. if ctr.Table != nil {
  1018. for id, cf := range ctr.Table.ColumnFamilies {
  1019. fams[id] = &columnFamily{
  1020. name: ctr.Parent + "/columnFamilies/" + id,
  1021. order: c,
  1022. gcRule: cf.GcRule,
  1023. }
  1024. c++
  1025. }
  1026. }
  1027. return &table{
  1028. families: fams,
  1029. counter: c,
  1030. rows: btree.New(btreeDegree),
  1031. }
  1032. }
  1033. func (t *table) validTimestamp(ts int64) bool {
  1034. if ts < minValidMilliSeconds || ts > maxValidMilliSeconds {
  1035. return false
  1036. }
  1037. // Assume millisecond granularity is required.
  1038. return ts%1000 == 0
  1039. }
  1040. func (t *table) columnFamilies() map[string]*columnFamily {
  1041. cp := make(map[string]*columnFamily)
  1042. t.mu.RLock()
  1043. for fam, cf := range t.families {
  1044. cp[fam] = cf
  1045. }
  1046. t.mu.RUnlock()
  1047. return cp
  1048. }
  1049. func (t *table) mutableRow(key string) *row {
  1050. bkey := btreeKey(key)
  1051. // Try fast path first.
  1052. t.mu.RLock()
  1053. i := t.rows.Get(bkey)
  1054. t.mu.RUnlock()
  1055. if i != nil {
  1056. return i.(*row)
  1057. }
  1058. // We probably need to create the row.
  1059. t.mu.Lock()
  1060. defer t.mu.Unlock()
  1061. i = t.rows.Get(bkey)
  1062. if i != nil {
  1063. return i.(*row)
  1064. }
  1065. r := newRow(key)
  1066. t.rows.ReplaceOrInsert(r)
  1067. return r
  1068. }
  1069. func (t *table) gc() {
  1070. // This method doesn't add or remove rows, so we only need a read lock for the table.
  1071. t.mu.RLock()
  1072. defer t.mu.RUnlock()
  1073. // Gather GC rules we'll apply.
  1074. rules := make(map[string]*btapb.GcRule) // keyed by "fam"
  1075. for fam, cf := range t.families {
  1076. if cf.gcRule != nil {
  1077. rules[fam] = cf.gcRule
  1078. }
  1079. }
  1080. if len(rules) == 0 {
  1081. return
  1082. }
  1083. t.rows.Ascend(func(i btree.Item) bool {
  1084. r := i.(*row)
  1085. r.mu.Lock()
  1086. r.gc(rules)
  1087. r.mu.Unlock()
  1088. return true
  1089. })
  1090. }
  1091. type byRowKey []*row
  1092. func (b byRowKey) Len() int { return len(b) }
  1093. func (b byRowKey) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
  1094. func (b byRowKey) Less(i, j int) bool { return b[i].key < b[j].key }
  1095. type row struct {
  1096. key string
  1097. mu sync.Mutex
  1098. families map[string]*family // keyed by family name
  1099. }
  1100. func newRow(key string) *row {
  1101. return &row{
  1102. key: key,
  1103. families: make(map[string]*family),
  1104. }
  1105. }
  1106. // copy returns a copy of the row.
  1107. // Cell values are aliased.
  1108. // r.mu should be held.
  1109. func (r *row) copy() *row {
  1110. nr := newRow(r.key)
  1111. for _, fam := range r.families {
  1112. nr.families[fam.name] = &family{
  1113. name: fam.name,
  1114. order: fam.order,
  1115. colNames: fam.colNames,
  1116. cells: make(map[string][]cell),
  1117. }
  1118. for col, cs := range fam.cells {
  1119. // Copy the []cell slice, but not the []byte inside each cell.
  1120. nr.families[fam.name].cells[col] = append([]cell(nil), cs...)
  1121. }
  1122. }
  1123. return nr
  1124. }
  1125. // isEmpty returns true if a row doesn't contain any cell
  1126. func (r *row) isEmpty() bool {
  1127. for _, fam := range r.families {
  1128. for _, cs := range fam.cells {
  1129. if len(cs) > 0 {
  1130. return false
  1131. }
  1132. }
  1133. }
  1134. return true
  1135. }
  1136. // sortedFamilies returns a column family set
  1137. // sorted in ascending creation order in a row.
  1138. func (r *row) sortedFamilies() []*family {
  1139. var families []*family
  1140. for _, fam := range r.families {
  1141. families = append(families, fam)
  1142. }
  1143. sort.Sort(byCreationOrder(families))
  1144. return families
  1145. }
  1146. func (r *row) getOrCreateFamily(name string, order uint64) *family {
  1147. if _, ok := r.families[name]; !ok {
  1148. r.families[name] = &family{
  1149. name: name,
  1150. order: order,
  1151. cells: make(map[string][]cell),
  1152. }
  1153. }
  1154. return r.families[name]
  1155. }
  1156. // gc applies the given GC rules to the row.
  1157. // r.mu should be held.
  1158. func (r *row) gc(rules map[string]*btapb.GcRule) {
  1159. for _, fam := range r.families {
  1160. rule, ok := rules[fam.name]
  1161. if !ok {
  1162. continue
  1163. }
  1164. for col, cs := range fam.cells {
  1165. r.families[fam.name].cells[col] = applyGC(cs, rule)
  1166. }
  1167. }
  1168. }
  1169. // size returns the total size of all cell values in the row.
  1170. func (r *row) size() int {
  1171. size := 0
  1172. for _, fam := range r.families {
  1173. for _, cells := range fam.cells {
  1174. for _, cell := range cells {
  1175. size += len(cell.value)
  1176. }
  1177. }
  1178. }
  1179. return size
  1180. }
  1181. // Less implements btree.Less.
  1182. func (r *row) Less(i btree.Item) bool {
  1183. return r.key < i.(*row).key
  1184. }
  1185. // btreeKey returns a row for use as a key into the BTree.
  1186. func btreeKey(s string) *row { return &row{key: s} }
  1187. func (r *row) String() string {
  1188. return r.key
  1189. }
  1190. var gcTypeWarn sync.Once
  1191. // applyGC applies the given GC rule to the cells.
  1192. func applyGC(cells []cell, rule *btapb.GcRule) []cell {
  1193. switch rule := rule.Rule.(type) {
  1194. default:
  1195. // TODO(dsymonds): Support GcRule_Intersection_
  1196. gcTypeWarn.Do(func() {
  1197. log.Printf("Unsupported GC rule type %T", rule)
  1198. })
  1199. case *btapb.GcRule_Union_:
  1200. for _, sub := range rule.Union.Rules {
  1201. cells = applyGC(cells, sub)
  1202. }
  1203. return cells
  1204. case *btapb.GcRule_MaxAge:
  1205. // Timestamps are in microseconds.
  1206. cutoff := time.Now().UnixNano() / 1e3
  1207. cutoff -= rule.MaxAge.Seconds * 1e6
  1208. cutoff -= int64(rule.MaxAge.Nanos) / 1e3
  1209. // The slice of cells in in descending timestamp order.
  1210. // This sort.Search will return the index of the first cell whose timestamp is chronologically before the cutoff.
  1211. si := sort.Search(len(cells), func(i int) bool { return cells[i].ts < cutoff })
  1212. if si < len(cells) {
  1213. log.Printf("bttest: GC MaxAge(%v) deleted %d cells.", rule.MaxAge, len(cells)-si)
  1214. }
  1215. return cells[:si]
  1216. case *btapb.GcRule_MaxNumVersions:
  1217. n := int(rule.MaxNumVersions)
  1218. if len(cells) > n {
  1219. cells = cells[:n]
  1220. }
  1221. return cells
  1222. }
  1223. return cells
  1224. }
  1225. type family struct {
  1226. name string // Column family name
  1227. order uint64 // Creation order of column family
  1228. colNames []string // Column names are sorted in lexicographical ascending order
  1229. cells map[string][]cell // Keyed by column name; cells are in descending timestamp order
  1230. }
  1231. type byCreationOrder []*family
  1232. func (b byCreationOrder) Len() int { return len(b) }
  1233. func (b byCreationOrder) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
  1234. func (b byCreationOrder) Less(i, j int) bool { return b[i].order < b[j].order }
  1235. // cellsByColumn adds the column name to colNames set if it does not exist
  1236. // and returns all cells within a column
  1237. func (f *family) cellsByColumn(name string) []cell {
  1238. if _, ok := f.cells[name]; !ok {
  1239. f.colNames = append(f.colNames, name)
  1240. sort.Strings(f.colNames)
  1241. }
  1242. return f.cells[name]
  1243. }
  1244. type cell struct {
  1245. ts int64
  1246. value []byte
  1247. labels []string
  1248. }
  1249. type byDescTS []cell
  1250. func (b byDescTS) Len() int { return len(b) }
  1251. func (b byDescTS) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
  1252. func (b byDescTS) Less(i, j int) bool { return b[i].ts > b[j].ts }
  1253. type columnFamily struct {
  1254. name string
  1255. order uint64 // Creation order of column family
  1256. gcRule *btapb.GcRule
  1257. }
  1258. func (c *columnFamily) proto() *btapb.ColumnFamily {
  1259. return &btapb.ColumnFamily{
  1260. GcRule: c.gcRule,
  1261. }
  1262. }
  1263. func toColumnFamilies(families map[string]*columnFamily) map[string]*btapb.ColumnFamily {
  1264. fs := make(map[string]*btapb.ColumnFamily)
  1265. for k, v := range families {
  1266. fs[k] = v.proto()
  1267. }
  1268. return fs
  1269. }