github.com/zuoyebang/bitalosdb@v1.1.1-0.20240516111551-79a8c4d8ce20/bitree/bdb/compact.go (about)

     1  // Copyright 2021 The Bitalosdb author(hustxrb@163.com) and other contributors.
     2  //
     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  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package bdb
    16  
    17  func Compact(dst, src *DB, ciMaxSize int64) error {
    18  	var size int64
    19  	tx, err := dst.Begin(true)
    20  	if err != nil {
    21  		return err
    22  	}
    23  	defer func() {
    24  		if err != nil {
    25  			_ = tx.Rollback()
    26  		}
    27  	}()
    28  
    29  	if err = walk(src, func(keys [][]byte, k, v []byte, seq uint64) error {
    30  		sz := int64(len(k) + len(v))
    31  		if size+sz > ciMaxSize && ciMaxSize != 0 {
    32  			if err := tx.Commit(); err != nil {
    33  				return err
    34  			}
    35  
    36  			tx, err = dst.Begin(true)
    37  			if err != nil {
    38  				return err
    39  			}
    40  			size = 0
    41  		}
    42  		size += sz
    43  
    44  		nk := len(keys)
    45  		if nk == 0 {
    46  			bkt, err := tx.CreateBucket(k)
    47  			if err != nil {
    48  				return err
    49  			}
    50  			if err := bkt.SetSequence(seq); err != nil {
    51  				return err
    52  			}
    53  			return nil
    54  		}
    55  
    56  		b := tx.Bucket(keys[0])
    57  		if nk > 1 {
    58  			for _, k := range keys[1:] {
    59  				b = b.Bucket(k)
    60  			}
    61  		}
    62  
    63  		b.FillPercent = 1.0
    64  
    65  		if v == nil {
    66  			bkt, err := b.CreateBucket(k)
    67  			if err != nil {
    68  				return err
    69  			}
    70  			if err := bkt.SetSequence(seq); err != nil {
    71  				return err
    72  			}
    73  			return nil
    74  		}
    75  
    76  		return b.Put(k, v)
    77  	}); err != nil {
    78  		return err
    79  	}
    80  
    81  	return tx.Commit()
    82  }
    83  
    84  type walkFunc func(keys [][]byte, k, v []byte, seq uint64) error
    85  
    86  func walk(db *DB, walkFn walkFunc) error {
    87  	return db.View(func(tx *Tx) error {
    88  		return tx.ForEach(func(name []byte, b *Bucket) error {
    89  			return walkBucket(b, nil, name, nil, b.Sequence(), walkFn)
    90  		})
    91  	})
    92  }
    93  
    94  func walkBucket(b *Bucket, keypath [][]byte, k, v []byte, seq uint64, fn walkFunc) error {
    95  	if err := fn(keypath, k, v, seq); err != nil {
    96  		return err
    97  	}
    98  
    99  	if v != nil {
   100  		return nil
   101  	}
   102  
   103  	keypath = append(keypath, k)
   104  	return b.ForEach(func(k, v []byte) error {
   105  		if v == nil {
   106  			bkt := b.Bucket(k)
   107  			return walkBucket(bkt, keypath, k, nil, bkt.Sequence(), fn)
   108  		}
   109  		return walkBucket(b, keypath, k, v, b.Sequence(), fn)
   110  	})
   111  }