codec_extension.go 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. // Copyright 2019 The Go Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package impl
  5. import (
  6. "sync"
  7. "sync/atomic"
  8. "google.golang.org/protobuf/internal/encoding/wire"
  9. pref "google.golang.org/protobuf/reflect/protoreflect"
  10. )
  11. type extensionFieldInfo struct {
  12. wiretag uint64
  13. tagsize int
  14. unmarshalNeedsValue bool
  15. funcs ifaceCoderFuncs
  16. }
  17. func (mi *MessageInfo) extensionFieldInfo(xt pref.ExtensionType) *extensionFieldInfo {
  18. // As of this time (Go 1.12, linux/amd64), an RWMutex benchmarks as faster
  19. // than a sync.Map.
  20. mi.extensionFieldInfosMu.RLock()
  21. e, ok := mi.extensionFieldInfos[xt]
  22. mi.extensionFieldInfosMu.RUnlock()
  23. if ok {
  24. return e
  25. }
  26. wiretag := wire.EncodeTag(xt.Number(), wireTypes[xt.Kind()])
  27. e = &extensionFieldInfo{
  28. wiretag: wiretag,
  29. tagsize: wire.SizeVarint(wiretag),
  30. funcs: encoderFuncsForValue(xt, xt.GoType()),
  31. }
  32. // Does the unmarshal function need a value passed to it?
  33. // This is true for composite types, where we pass in a message, list, or map to fill in,
  34. // and for enums, where we pass in a prototype value to specify the concrete enum type.
  35. switch xt.Kind() {
  36. case pref.MessageKind, pref.GroupKind, pref.EnumKind:
  37. e.unmarshalNeedsValue = true
  38. default:
  39. if xt.Cardinality() == pref.Repeated {
  40. e.unmarshalNeedsValue = true
  41. }
  42. }
  43. mi.extensionFieldInfosMu.Lock()
  44. if mi.extensionFieldInfos == nil {
  45. mi.extensionFieldInfos = make(map[pref.ExtensionType]*extensionFieldInfo)
  46. }
  47. mi.extensionFieldInfos[xt] = e
  48. mi.extensionFieldInfosMu.Unlock()
  49. return e
  50. }
  51. type ExtensionField struct {
  52. typ pref.ExtensionType
  53. // value is either the value of GetValue,
  54. // or a *lazyExtensionValue that then returns the value of GetValue.
  55. value interface{} // TODO: switch to protoreflect.Value
  56. }
  57. func (f ExtensionField) HasType() bool {
  58. return f.typ != nil
  59. }
  60. func (f ExtensionField) GetType() pref.ExtensionType {
  61. return f.typ
  62. }
  63. func (f *ExtensionField) SetType(t pref.ExtensionType) {
  64. f.typ = t
  65. }
  66. // HasValue reports whether a value is set for the extension field.
  67. // This may be called concurrently.
  68. func (f ExtensionField) HasValue() bool {
  69. return f.value != nil
  70. }
  71. // GetValue returns the concrete value for the extension field.
  72. // Let the type of Desc.ExtensionType be the "API type" and
  73. // the type of GetValue be the "storage type".
  74. // The API type and storage type are the same except:
  75. // * for scalars (except []byte), where the API type uses *T,
  76. // while the storage type uses T.
  77. // * for repeated fields, where the API type uses []T,
  78. // while the storage type uses *[]T.
  79. //
  80. // The reason for the divergence is so that the storage type more naturally
  81. // matches what is expected of when retrieving the values through the
  82. // protobuf reflection APIs.
  83. //
  84. // GetValue is only populated if Desc is also populated.
  85. // This may be called concurrently.
  86. //
  87. // TODO: switch interface{} to protoreflect.Value
  88. func (f ExtensionField) GetValue() interface{} {
  89. if f, ok := f.value.(*lazyExtensionValue); ok {
  90. return f.GetValue()
  91. }
  92. return f.value
  93. }
  94. // SetEagerValue sets the current value of the extension.
  95. // This must not be called concurrently.
  96. func (f *ExtensionField) SetEagerValue(v interface{}) {
  97. f.value = v
  98. }
  99. // SetLazyValue sets a value that is to be lazily evaluated upon first use.
  100. // The returned value must not be nil.
  101. // This must not be called concurrently.
  102. func (f *ExtensionField) SetLazyValue(v func() interface{}) {
  103. f.value = &lazyExtensionValue{value: v}
  104. }
  105. type lazyExtensionValue struct {
  106. once uint32 // atomically set if value is valid
  107. mu sync.Mutex // protects value
  108. value interface{} // either the value itself or a func() interface{}
  109. }
  110. func (v *lazyExtensionValue) GetValue() interface{} {
  111. if atomic.LoadUint32(&v.once) == 0 {
  112. v.mu.Lock()
  113. if f, ok := v.value.(func() interface{}); ok {
  114. v.value = f()
  115. }
  116. atomic.StoreUint32(&v.once, 1)
  117. v.mu.Unlock()
  118. }
  119. return v.value
  120. }