github.com/weaviate/weaviate@v1.24.6/entities/diskio/metered_reader.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package diskio
    13  
    14  import (
    15  	"io"
    16  	"time"
    17  )
    18  
    19  type MeteredReaderCallback func(read int64, nanoseconds int64)
    20  
    21  type MeteredReader struct {
    22  	r  io.Reader
    23  	cb MeteredReaderCallback
    24  }
    25  
    26  // Read passes the read through to the underlying reader. On a successful read,
    27  // it will trigger the attached callback and provide it with metrics. If no
    28  // callback is set, it will ignore it.
    29  func (m *MeteredReader) Read(p []byte) (n int, err error) {
    30  	start := time.Now()
    31  	n, err = m.r.Read(p)
    32  	took := time.Since(start).Nanoseconds()
    33  	if err != nil {
    34  		return
    35  	}
    36  
    37  	if m.cb != nil {
    38  		m.cb(int64(n), took)
    39  	}
    40  
    41  	return
    42  }
    43  
    44  func NewMeteredReader(r io.Reader, cb MeteredReaderCallback) *MeteredReader {
    45  	return &MeteredReader{r: r, cb: cb}
    46  }