@@ -12,113 +12,145 @@ import (
1212 "github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
1313 "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
1414 "github.com/cockroachdb/cockroach/pkg/sql/catalog/nstree"
15- "github.com/cockroachdb/cockroach/pkg/sql/catalog/zone"
1615)
1716
18- type uncommittedComments struct {
19- uncommitted map [catalogkeys.CommentKey ]string
20- cachedKeys map [catalogkeys.CommentKey ]struct {}
21- }
17+ // uncommittedMetadata tracks uncommitted comments and zone configs using
18+ // an nstree.MutableCatalog for storage, plus separate sets for tracking
19+ // entries that have been explicitly marked as absent.
20+ type uncommittedMetadata struct {
21+ // mc stores actual uncommitted values (comments and zone configs).
22+ // This handles both storage and memory tracking via its built-in byteSize.
23+ mc nstree.MutableCatalog
2224
23- func makeUncommittedComments () uncommittedComments {
24- return uncommittedComments {}
25- }
25+ // absentCommentKeys tracks comment keys that have been looked up and
26+ // found to be absent (deleted or never existed). This is separate from
27+ // mc because MutableCatalog only stores actual values.
28+ absentCommentKeys map [catalogkeys.CommentKey ]struct {}
2629
27- func (uc * uncommittedComments ) reset () {
28- * uc = uncommittedComments {}
30+ // absentZoneConfigIDs tracks descriptor IDs whose zone configs have
31+ // been looked up and found to be absent.
32+ absentZoneConfigIDs map [descpb.ID ]struct {}
2933}
3034
31- func (uc * uncommittedComments ) lazyInitMaps () {
32- if uc .uncommitted == nil {
33- uc .uncommitted = make (map [catalogkeys.CommentKey ]string )
34- uc .cachedKeys = make (map [catalogkeys.CommentKey ]struct {})
35- }
35+ func makeUncommittedMetadata () uncommittedMetadata {
36+ return uncommittedMetadata {}
3637}
3738
38- func (uc * uncommittedComments ) getUncommitted (
39- key catalogkeys.CommentKey ,
40- ) (cmt string , hasCmt bool , cached bool ) {
41- if _ , ok := uc .cachedKeys [key ]; ! ok {
42- return "" , false , false
43- }
44-
45- cmt , hasCmt = uc .uncommitted [key ]
46- return cmt , hasCmt , true
39+ func (um * uncommittedMetadata ) reset () {
40+ um .mc .Clear ()
41+ um .absentCommentKeys = nil
42+ um .absentZoneConfigIDs = nil
4743}
4844
49- // markNoComment lets the cache know that the comment for this key is dropped.
50- func (uc * uncommittedComments ) markNoComment (key catalogkeys.CommentKey ) {
51- uc .lazyInitMaps ()
52- delete (uc .uncommitted , key )
53- uc .cachedKeys [key ] = struct {}{}
54- }
45+ // Comment methods
5546
56- func (uc * uncommittedComments ) markTableDeleted (tblID descpb.ID ) {
57- // NOTE: lazyInitMaps() not needed, maps can remain nil.
58- var keysToDel []catalogkeys.CommentKey
59- for k := range uc .uncommitted {
60- if k .ObjectID == uint32 (tblID ) {
61- keysToDel = append (keysToDel , k )
47+ // getUncommittedComment returns the uncommitted comment for the given key.
48+ // It returns (comment, hasComment, cached) where:
49+ // - cached=false means the key hasn't been looked up yet
50+ // - cached=true, hasComment=false means the key was looked up and is absent
51+ // - cached=true, hasComment=true means the key has a value
52+ func (um * uncommittedMetadata ) getUncommittedComment (
53+ key catalogkeys.CommentKey ,
54+ ) (cmt string , hasCmt bool , cached bool ) {
55+ if um .absentCommentKeys != nil {
56+ if _ , absent := um .absentCommentKeys [key ]; absent {
57+ return "" , false , true
6258 }
6359 }
64- for _ , k := range keysToDel {
65- delete ( uc . uncommitted , k )
60+ if cmt , found := um . mc . LookupComment ( key ); found {
61+ return cmt , true , true
6662 }
63+ return "" , false , false
6764}
6865
69- func (uc * uncommittedComments ) upsert (key catalogkeys.CommentKey , cmt string ) {
70- uc .lazyInitMaps ()
71- uc .cachedKeys [key ] = struct {}{}
72- uc .uncommitted [key ] = cmt
73- }
74-
75- func (uc * uncommittedComments ) addAllToCatalog (mc nstree.MutableCatalog ) error {
76- for ck , cmt := range uc .uncommitted {
77- if err := mc .UpsertComment (ck , cmt ); err != nil {
78- return err
79- }
66+ // markNoComment marks the comment for this key as absent (deleted or never
67+ // existed). This adds it to the cache so subsequent lookups know not to
68+ // query storage.
69+ func (um * uncommittedMetadata ) markNoComment (key catalogkeys.CommentKey ) {
70+ um .mc .DeleteComment (key )
71+ if um .absentCommentKeys == nil {
72+ um .absentCommentKeys = make (map [catalogkeys.CommentKey ]struct {})
8073 }
81- return nil
74+ um . absentCommentKeys [ key ] = struct {}{}
8275}
8376
84- type uncommittedZoneConfigs struct {
85- uncommitted map [descpb.ID ]catalog.ZoneConfig
86- cachedDescs map [descpb.ID ]struct {}
77+ // upsertComment adds or updates a comment in the uncommitted layer.
78+ func (um * uncommittedMetadata ) upsertComment (key catalogkeys.CommentKey , cmt string ) error {
79+ delete (um .absentCommentKeys , key )
80+ return um .mc .UpsertComment (key , cmt )
8781}
8882
89- func makeUncommittedZoneConfigs () uncommittedZoneConfigs {
90- return uncommittedZoneConfigs {}
83+ // markTableCommentsDeleted removes all comments for the given table from the
84+ // uncommitted layer. Note: this only affects comments that were previously
85+ // added to the uncommitted layer; it does not mark stored comments as deleted.
86+ func (um * uncommittedMetadata ) markTableCommentsDeleted (tblID descpb.ID ) {
87+ // Collect keys to delete from MutableCatalog.
88+ var keysToDelete []catalogkeys.CommentKey
89+ _ = um .mc .ForEachComment (func (key catalogkeys.CommentKey , _ string ) error {
90+ if descpb .ID (key .ObjectID ) == tblID {
91+ keysToDelete = append (keysToDelete , key )
92+ }
93+ return nil
94+ })
95+ for _ , key := range keysToDelete {
96+ um .mc .DeleteComment (key )
97+ }
9198}
9299
93- func (uc * uncommittedZoneConfigs ) reset () {
94- * uc = uncommittedZoneConfigs {}
100+ // addAllCommentsToCatalog copies all uncommitted comments to the target catalog.
101+ func (um * uncommittedMetadata ) addAllCommentsToCatalog (target * nstree.MutableCatalog ) error {
102+ return um .mc .ForEachComment (func (key catalogkeys.CommentKey , cmt string ) error {
103+ return target .UpsertComment (key , cmt )
104+ })
95105}
96106
97- func (uc * uncommittedZoneConfigs ) lazyInitMaps () {
98- if uc .uncommitted == nil {
99- uc .uncommitted = make (map [descpb.ID ]catalog.ZoneConfig )
100- uc .cachedDescs = make (map [descpb.ID ]struct {})
107+ // isCommentCached returns true if the comment key is in the cache (either as
108+ // a value or marked absent).
109+ func (um * uncommittedMetadata ) isCommentCached (key catalogkeys.CommentKey ) bool {
110+ if um .absentCommentKeys != nil {
111+ if _ , absent := um .absentCommentKeys [key ]; absent {
112+ return true
113+ }
101114 }
115+ _ , found := um .mc .LookupComment (key )
116+ return found
102117}
103118
104- func (uc * uncommittedZoneConfigs ) getUncommitted (
119+ // Zone config methods
120+
121+ // getUncommittedZoneConfig returns the uncommitted zone config for the given ID.
122+ // It returns (zoneConfig, cached) where:
123+ // - cached=false means the ID hasn't been looked up yet
124+ // - cached=true, zoneConfig=nil means the ID was looked up and is absent
125+ // - cached=true, zoneConfig!=nil means the ID has a value
126+ func (um * uncommittedMetadata ) getUncommittedZoneConfig (
105127 id descpb.ID ,
106128) (zc catalog.ZoneConfig , cached bool ) {
107- if _ , ok := uc .cachedDescs [id ]; ! ok {
108- return nil , false
129+ if um .absentZoneConfigIDs != nil {
130+ if _ , absent := um .absentZoneConfigIDs [id ]; absent {
131+ return nil , true
132+ }
133+ }
134+ if zc := um .mc .LookupZoneConfig (id ); zc != nil {
135+ return zc , true
109136 }
110- return uc . uncommitted [ id ], true
137+ return nil , false
111138}
112139
113- func (uc * uncommittedZoneConfigs ) markNoZoneConfig (id descpb.ID ) {
114- uc .lazyInitMaps ()
115- delete (uc .uncommitted , id )
116- uc .cachedDescs [id ] = struct {}{}
140+ // markNoZoneConfig marks the zone config for this ID as absent (deleted or
141+ // never existed). This adds it to the cache so subsequent lookups know not
142+ // to query storage.
143+ func (um * uncommittedMetadata ) markNoZoneConfig (id descpb.ID ) {
144+ um .mc .DeleteZoneConfig (id )
145+ if um .absentZoneConfigIDs == nil {
146+ um .absentZoneConfigIDs = make (map [descpb.ID ]struct {})
147+ }
148+ um .absentZoneConfigIDs [id ] = struct {}{}
117149}
118150
119- func ( uc * uncommittedZoneConfigs ) upsert ( id descpb. ID , zc * zonepb. ZoneConfig ) error {
120- uc . lazyInitMaps ()
121- uc . cachedDescs [ id ] = struct {}{}
151+ // upsertZoneConfig adds or updates a zone config in the uncommitted layer.
152+ func ( um * uncommittedMetadata ) upsertZoneConfig ( id descpb. ID , zc * zonepb. ZoneConfig ) error {
153+ delete ( um . absentZoneConfigIDs , id )
122154 var val roachpb.Value
123155 if err := val .SetProto (zc ); err != nil {
124156 return err
@@ -127,12 +159,26 @@ func (uc *uncommittedZoneConfigs) upsert(id descpb.ID, zc *zonepb.ZoneConfig) er
127159 if err != nil {
128160 return err
129161 }
130- uc . uncommitted [ id ] = zone . NewZoneConfigWithRawBytes ( zc , rawBytes )
162+ um . mc . UpsertZoneConfig ( id , zc , rawBytes )
131163 return nil
132164}
133165
134- func (uc * uncommittedZoneConfigs ) addAllToCatalog (mc nstree.MutableCatalog ) {
135- for id , zc := range uc .uncommitted {
136- mc .UpsertZoneConfig (id , zc .ZoneConfigProto (), zc .GetRawBytesInStorage ())
166+ // addAllZoneConfigsToCatalog copies all uncommitted zone configs to the target
167+ // catalog.
168+ func (um * uncommittedMetadata ) addAllZoneConfigsToCatalog (target * nstree.MutableCatalog ) {
169+ _ = um .mc .ForEachZoneConfig (func (id descpb.ID , zc catalog.ZoneConfig ) error {
170+ target .UpsertZoneConfig (id , zc .ZoneConfigProto (), zc .GetRawBytesInStorage ())
171+ return nil
172+ })
173+ }
174+
175+ // isZoneConfigCached returns true if the zone config ID is in the cache
176+ // (either as a value or marked absent).
177+ func (um * uncommittedMetadata ) isZoneConfigCached (id descpb.ID ) bool {
178+ if um .absentZoneConfigIDs != nil {
179+ if _ , absent := um .absentZoneConfigIDs [id ]; absent {
180+ return true
181+ }
137182 }
183+ return um .mc .LookupZoneConfig (id ) != nil
138184}
0 commit comments