github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/storage/stores/shipper/index/indexfile/indexfile.go (about)

     1  package indexfile
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"path"
     9  	"path/filepath"
    10  	"sync"
    11  	"time"
    12  
    13  	"github.com/go-kit/log/level"
    14  	"go.etcd.io/bbolt"
    15  
    16  	"github.com/grafana/loki/pkg/storage/chunk/client/local"
    17  	"github.com/grafana/loki/pkg/storage/stores/indexshipper/index"
    18  	series_index "github.com/grafana/loki/pkg/storage/stores/series/index"
    19  	"github.com/grafana/loki/pkg/storage/stores/shipper/util"
    20  	util_log "github.com/grafana/loki/pkg/util/log"
    21  )
    22  
    23  const TempFileSuffix = ".temp"
    24  
    25  type IndexFile struct {
    26  	boltDB *bbolt.DB
    27  	name   string
    28  
    29  	path             string
    30  	issuedReaders    []*os.File
    31  	issuedReadersMtx sync.Mutex
    32  }
    33  
    34  func OpenIndexFile(path string) (index.Index, error) {
    35  	boltdbFile, err := util.SafeOpenBoltdbFile(path)
    36  	if err != nil {
    37  		return nil, err
    38  	}
    39  
    40  	return &IndexFile{
    41  		boltDB: boltdbFile,
    42  		name:   filepath.Base(path),
    43  		path:   path,
    44  	}, nil
    45  }
    46  
    47  func BoltDBToIndexFile(boltdbFile *bbolt.DB, name string) index.Index {
    48  	return &IndexFile{
    49  		boltDB: boltdbFile,
    50  		name:   name,
    51  		path:   boltdbFile.Path(),
    52  	}
    53  }
    54  
    55  func (i *IndexFile) GetBoltDB() *bbolt.DB {
    56  	return i.boltDB
    57  }
    58  
    59  func (i *IndexFile) Path() string {
    60  	return i.path
    61  }
    62  
    63  func (i *IndexFile) Name() string {
    64  	return i.name
    65  }
    66  
    67  func (i *IndexFile) Reader() (io.ReadSeeker, error) {
    68  	filePath := path.Join(filepath.Dir(i.Path()), fmt.Sprintf("%d%s", time.Now().UnixNano(), TempFileSuffix))
    69  	f, err := os.Create(filePath)
    70  	if err != nil {
    71  		return nil, err
    72  	}
    73  
    74  	err = i.boltDB.View(func(tx *bbolt.Tx) (err error) {
    75  		_, err = tx.WriteTo(f)
    76  		return
    77  	})
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  
    82  	// flush the file to disk and seek the file to the beginning.
    83  	if err := f.Sync(); err != nil {
    84  		return nil, err
    85  	}
    86  
    87  	if _, err := f.Seek(0, 0); err != nil {
    88  		return nil, err
    89  	}
    90  
    91  	i.issuedReadersMtx.Lock()
    92  	defer i.issuedReadersMtx.Unlock()
    93  	i.issuedReaders = append(i.issuedReaders, f)
    94  
    95  	return f, nil
    96  }
    97  
    98  func (i *IndexFile) Close() error {
    99  	i.issuedReadersMtx.Lock()
   100  	defer i.issuedReadersMtx.Unlock()
   101  
   102  	// cleanup all the issued readers
   103  	for _, f := range i.issuedReaders {
   104  		if err := f.Close(); err != nil {
   105  			level.Error(util_log.Logger).Log("msg", "failed to close temp file", "path", f.Name(), "err", err)
   106  		}
   107  
   108  		if err := os.Remove(f.Name()); err != nil {
   109  			level.Error(util_log.Logger).Log("msg", "failed to remove temp file", "path", f.Name(), "err", err)
   110  		}
   111  	}
   112  
   113  	return i.boltDB.Close()
   114  }
   115  
   116  func QueryBoltDB(ctx context.Context, db *bbolt.DB, userID []byte, queries []series_index.Query, callback series_index.QueryPagesCallback) error {
   117  	return db.View(func(tx *bbolt.Tx) error {
   118  		bucket := tx.Bucket(userID)
   119  		if bucket == nil {
   120  			bucket = tx.Bucket(local.IndexBucketName)
   121  			if bucket == nil {
   122  				return nil
   123  			}
   124  		}
   125  
   126  		for _, query := range queries {
   127  			if err := local.QueryWithCursor(ctx, bucket.Cursor(), query, callback); err != nil {
   128  				return err
   129  			}
   130  		}
   131  		return nil
   132  	})
   133  }