github.com/jxskiss/gopkg/v2@v2.14.9-0.20240514120614-899f3e7952b4/utils/sqlutil/types_binary.go (about)

     1  package sqlutil
     2  
     3  import (
     4  	"database/sql/driver"
     5  	"fmt"
     6  	"sync"
     7  	"sync/atomic"
     8  	"unsafe"
     9  )
    10  
    11  // NewLazyBinary creates a new lazy binary wrapper, delaying the
    12  // unmarshalling work until it is first needed.
    13  func NewLazyBinary(raw []byte) LazyBinary {
    14  	return LazyBinary{raw: raw}
    15  }
    16  
    17  // LazyBinary is a lazy wrapper around a binary value, where the underlying
    18  // object will be unmarshalled the first time it is needed and cached.
    19  // It implements sql/driver.Valuer and sql.Scanner.
    20  //
    21  // LazyBinary provides same concurrency safety as []byte, it's safe for
    22  // concurrent read, but not safe for concurrent write or read/write.
    23  //
    24  // See types_test.go for example usage.
    25  type LazyBinary struct {
    26  	raw []byte
    27  	obj unsafe.Pointer // *lazyobj
    28  }
    29  
    30  type lazyobj struct {
    31  	mu   sync.Mutex
    32  	data any
    33  	err  error
    34  }
    35  
    36  // Value implements driver.Valuer interface.
    37  func (p LazyBinary) Value() (driver.Value, error) {
    38  	return p.raw, nil
    39  }
    40  
    41  // Scan implements sql.Scanner interface.
    42  func (p *LazyBinary) Scan(src any) error {
    43  	if src != nil {
    44  		// NOTE
    45  		// We MUST copy the src byte slice here, database/sql.Scanner says:
    46  		//
    47  		// Reference types such as []byte are only valid until the next call to Scan
    48  		// and should not be retained. Their underlying memory is owned by the driver.
    49  		// If retention is necessary, copy their values before the next call to Scan.
    50  
    51  		var b []byte
    52  		switch tmp := src.(type) {
    53  		case string:
    54  			b = []byte(tmp)
    55  		case []byte:
    56  			b = make([]byte, len(tmp))
    57  			copy(b, tmp)
    58  		default:
    59  			return fmt.Errorf("sqlutil.LazyBinary.Scan: want string/[]byte but got %T", src)
    60  		}
    61  		p.raw = b
    62  		atomic.StorePointer(&p.obj, nil)
    63  	}
    64  	return nil
    65  }
    66  
    67  // Unmarshaler is a function which unmarshalls data from a byte slice.
    68  type Unmarshaler func([]byte) (any, error)
    69  
    70  // GetBytes returns the underlying byte slice.
    71  func (p *LazyBinary) GetBytes() []byte {
    72  	return p.raw
    73  }
    74  
    75  // Get returns the underlying data wrapped by the LazyBinary wrapper,
    76  // if the data has not been unmarshalled, it will be unmarshalled using
    77  // the provided unmarshalFunc.
    78  // The unmarshalling work will do only once, the result data and error
    79  // will be cached and reused for further calling.
    80  func (p *LazyBinary) Get(unmarshalFunc Unmarshaler) (any, error) {
    81  	obj, created := p.getobj()
    82  	defer obj.mu.Unlock()
    83  	if created {
    84  		obj.data, obj.err = unmarshalFunc(p.raw)
    85  		return obj.data, obj.err
    86  	}
    87  	obj.mu.Lock()
    88  	return obj.data, obj.err
    89  }
    90  
    91  // Set sets the data and marshaled bytes to the LazyBinary wrapper.
    92  // If the param data is nil, the underlying cache will be removed.
    93  func (p *LazyBinary) Set(b []byte, data any) {
    94  	p.raw = b
    95  	if data == nil {
    96  		atomic.StorePointer(&p.obj, nil)
    97  		return
    98  	}
    99  	obj, created := p.getobj()
   100  	if !created {
   101  		obj.mu.Lock()
   102  	}
   103  	obj.data = data
   104  	obj.err = nil
   105  	obj.mu.Unlock()
   106  }
   107  
   108  func (p *LazyBinary) getobj() (*lazyobj, bool) {
   109  	ptr := atomic.LoadPointer(&p.obj)
   110  	if ptr != nil {
   111  		return (*lazyobj)(ptr), false
   112  	}
   113  	tmp := &lazyobj{}
   114  	tmp.mu.Lock()
   115  	swapped := atomic.CompareAndSwapPointer(&p.obj, nil, unsafe.Pointer(tmp))
   116  	if swapped {
   117  		return tmp, true
   118  	}
   119  	ptr = atomic.LoadPointer(&p.obj)
   120  	return (*lazyobj)(ptr), false
   121  }