@@ -7,6 +7,7 @@ package changefeedccl
77
88import (
99 "bytes"
10+ "cmp"
1011 "context"
1112 "crypto/rand"
1213 "encoding/hex"
@@ -18,7 +19,7 @@ import (
1819 "sync/atomic"
1920 "time"
2021
21- "github.com/RaduBerinde/btree" // TODO(#144504): switch to the newer btree
22+ "github.com/RaduBerinde/btreemap"
2223 "github.com/cockroachdb/cockroach/pkg/base"
2324 "github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/changefeedbase"
2425 "github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/kvevent"
@@ -320,7 +321,7 @@ type cloudStorageSink struct {
320321 // These are fields to track information needed to output files based on the naming
321322 // convention described above. See comment on cloudStorageSink above for more details.
322323 fileID int64
323- files * btree. BTree // of *cloudStorageSinkFile
324+ files * btreemap. BTreeMap [ cloudStorageSinkKey , * cloudStorageSinkFile ]
324325
325326 timestampOracle timestampLowerBoundOracle
326327 jobSessionID string
@@ -416,7 +417,7 @@ func makeCloudStorageSink(
416417 sinkID : sinkID ,
417418 settings : settings ,
418419 targetMaxFileSize : targetMaxFileSize ,
419- files : btree .New ( 8 ),
420+ files : btreemap .New [ cloudStorageSinkKey , * cloudStorageSinkFile ]( 8 , keyCmp ),
420421 partitionFormat : defaultPartitionFormat ,
421422 timestampOracle : timestampOracle ,
422423 // TODO(dan,ajwerner): Use the jobs framework's session ID once that's available.
@@ -516,8 +517,7 @@ func (s *cloudStorageSink) getOrCreateFile(
516517) (* cloudStorageSinkFile , error ) {
517518 name , _ := s .topicNamer .Name (topic )
518519 key := cloudStorageSinkKey {name , int64 (topic .GetVersion ())}
519- if item := s .files .Get (key ); item != nil {
520- f := item .(* cloudStorageSinkFile )
520+ if _ , f , _ := s .files .Get (key ); f != nil {
521521 if eventMVCC .Less (f .oldestMVCC ) {
522522 f .oldestMVCC = eventMVCC
523523 }
@@ -537,7 +537,7 @@ func (s *cloudStorageSink) getOrCreateFile(
537537 }
538538 f .codec = codec
539539 }
540- s .files .ReplaceOrInsert (f )
540+ s .files .ReplaceOrInsert (f . cloudStorageSinkKey , f )
541541 return f , nil
542542}
543543
@@ -646,20 +646,17 @@ func (s *cloudStorageSink) EmitResolvedTimestamp(
646646// on cloudStorageSink)
647647func (s * cloudStorageSink ) flushTopicVersions (
648648 ctx context.Context , topic string , maxVersionToFlush int64 ,
649- ) ( err error ) {
649+ ) error {
650650 var toRemoveAlloc [2 ]int64 // generally avoid allocating
651651 toRemove := toRemoveAlloc [:0 ] // schemaIDs of flushed files
652652 gte := cloudStorageSinkKey {topic : topic }
653653 lt := cloudStorageSinkKey {topic : topic , schemaID : maxVersionToFlush + 1 }
654- s . files . AscendRange ( gte , lt , func ( i btree. Item ) ( wantMore bool ) {
655- f := i .( * cloudStorageSinkFile )
656- if err = s .flushFile (ctx , f ); err = =nil {
657- toRemove = append ( toRemove , f . schemaID )
654+
655+ for _ , f := range s . files . Ascend ( btreemap . GE ( gte ), btreemap . LT ( lt )) {
656+ if err : =s .flushFile (ctx , f ); err ! =nil {
657+ return err
658658 }
659- return err == nil
660- })
661- if err != nil {
662- return err
659+ toRemove = append (toRemove , f .schemaID )
663660 }
664661
665662 // Allow synchronization with the async flusher to happen.
@@ -672,8 +669,7 @@ func (s *cloudStorageSink) flushTopicVersions(
672669 // flushed files may not be removed from s.files. This is ok, since
673670 // the error will trigger the sink to be closed, and we will only use
674671 // s.files to ensure that the codecs are closed before deallocating it.
675- err = s .waitAsyncFlush (ctx )
676- if err != nil {
672+ if err := s .waitAsyncFlush (ctx ); err != nil {
677673 return err
678674 }
679675
@@ -682,7 +678,7 @@ func (s *cloudStorageSink) flushTopicVersions(
682678 for _ , v := range toRemove {
683679 s .files .Delete (cloudStorageSinkKey {topic : topic , schemaID : v })
684680 }
685- return err
681+ return nil
686682}
687683
688684// Flush implements the Sink interface.
@@ -693,13 +689,10 @@ func (s *cloudStorageSink) Flush(ctx context.Context) error {
693689
694690 s .metrics .recordFlushRequestCallback ()()
695691
696- var err error
697- s .files .Ascend (func (i btree.Item ) (wantMore bool ) {
698- err = s .flushFile (ctx , i .(* cloudStorageSinkFile ))
699- return err == nil
700- })
701- if err != nil {
702- return err
692+ for _ , f := range s .files .Ascend (btreemap .Min [cloudStorageSinkKey ](), btreemap .Max [cloudStorageSinkKey ]()) {
693+ if err := s .flushFile (ctx , f ); err != nil {
694+ return err
695+ }
703696 }
704697 // Allow synchronization with the async flusher to happen.
705698 if s .testingKnobs != nil && s .testingKnobs .AsyncFlushSync != nil {
@@ -711,8 +704,7 @@ func (s *cloudStorageSink) Flush(ctx context.Context) error {
711704 // flushed files may not be removed from s.files. This is ok, since
712705 // the error will trigger the sink to be closed, and we will only use
713706 // s.files to ensure that the codecs are closed before deallocating it.
714- err = s .waitAsyncFlush (ctx )
715- if err != nil {
707+ if err := s .waitAsyncFlush (ctx ); err != nil {
716708 return err
717709 }
718710 // Files need to be cleared after the flush completes, otherwise file resources
@@ -909,17 +901,15 @@ func (s *cloudStorageSink) closeAllCodecs() (err error) {
909901 // Codecs need to be closed because of the klauspost compression library implementation
910902 // details where it spins up go routines to perform compression in parallel.
911903 // Those go routines are cleaned up when the compression codec is closed.
912- s .files .Ascend (func (i btree.Item ) (wantMore bool ) {
913- f := i .(* cloudStorageSinkFile )
904+ for _ , f := range s .files .Ascend (btreemap .Min [cloudStorageSinkKey ](), btreemap .Max [cloudStorageSinkKey ]()) {
914905 if f .codec != nil {
915906 cErr := f .codec .Close ()
916907 f .codec = nil
917908 if err == nil {
918909 err = cErr
919910 }
920911 }
921- return true
922- })
912+ }
923913 return err
924914}
925915
@@ -943,22 +933,11 @@ type cloudStorageSinkKey struct {
943933 schemaID int64
944934}
945935
946- func (k cloudStorageSinkKey ) Less (other btree.Item ) bool {
947- switch other := other .(type ) {
948- case * cloudStorageSinkFile :
949- return keyLess (k , other .cloudStorageSinkKey )
950- case cloudStorageSinkKey :
951- return keyLess (k , other )
952- default :
953- panic (errors .Errorf ("unexpected item type %T" , other ))
954- }
955- }
956-
957- func keyLess (a , b cloudStorageSinkKey ) bool {
958- if a .topic == b .topic {
959- return a .schemaID < b .schemaID
936+ func keyCmp (a , b cloudStorageSinkKey ) int {
937+ if a .topic != b .topic {
938+ return cmp .Compare (a .topic , b .topic )
960939 }
961- return a . topic < b . topic
940+ return cmp . Compare ( a . schemaID , b . schemaID )
962941}
963942
964943// generateChangefeedSessionID generates a unique string that is used to
![craig[bot]](https://github.githubassets.com/images/gravatars/gravatar-user-420.png?size=40){kind=avatar}
0 commit comments