init.go (4604B)
1 // Copyright 2019 The Hugo Authors. All rights reserved.
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 // http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13
14 package lazy
15
16 import (
17 "context"
18 "sync"
19 "sync/atomic"
20 "time"
21
22 "errors"
23 )
24
25 // New creates a new empty Init.
26 func New() *Init {
27 return &Init{}
28 }
29
30 // Init holds a graph of lazily initialized dependencies.
31 type Init struct {
32 // Used in tests
33 initCount uint64
34
35 mu sync.Mutex
36
37 prev *Init
38 children []*Init
39
40 init onceMore
41 out any
42 err error
43 f func() (any, error)
44 }
45
46 // Add adds a func as a new child dependency.
47 func (ini *Init) Add(initFn func() (any, error)) *Init {
48 if ini == nil {
49 ini = New()
50 }
51 return ini.add(false, initFn)
52 }
53
54 // InitCount gets the number of this this Init has been initialized.
55 func (ini *Init) InitCount() int {
56 i := atomic.LoadUint64(&ini.initCount)
57 return int(i)
58 }
59
60 // AddWithTimeout is same as Add, but with a timeout that aborts initialization.
61 func (ini *Init) AddWithTimeout(timeout time.Duration, f func(ctx context.Context) (any, error)) *Init {
62 return ini.Add(func() (any, error) {
63 return ini.withTimeout(timeout, f)
64 })
65 }
66
67 // Branch creates a new dependency branch based on an existing and adds
68 // the given dependency as a child.
69 func (ini *Init) Branch(initFn func() (any, error)) *Init {
70 if ini == nil {
71 ini = New()
72 }
73 return ini.add(true, initFn)
74 }
75
76 // BranchdWithTimeout is same as Branch, but with a timeout.
77 func (ini *Init) BranchWithTimeout(timeout time.Duration, f func(ctx context.Context) (any, error)) *Init {
78 return ini.Branch(func() (any, error) {
79 return ini.withTimeout(timeout, f)
80 })
81 }
82
83 // Do initializes the entire dependency graph.
84 func (ini *Init) Do() (any, error) {
85 if ini == nil {
86 panic("init is nil")
87 }
88
89 ini.init.Do(func() {
90 atomic.AddUint64(&ini.initCount, 1)
91 prev := ini.prev
92 if prev != nil {
93 // A branch. Initialize the ancestors.
94 if prev.shouldInitialize() {
95 _, err := prev.Do()
96 if err != nil {
97 ini.err = err
98 return
99 }
100 } else if prev.inProgress() {
101 // Concurrent initialization. The following init func
102 // may depend on earlier state, so wait.
103 prev.wait()
104 }
105 }
106
107 if ini.f != nil {
108 ini.out, ini.err = ini.f()
109 }
110
111 for _, child := range ini.children {
112 if child.shouldInitialize() {
113 _, err := child.Do()
114 if err != nil {
115 ini.err = err
116 return
117 }
118 }
119 }
120 })
121
122 ini.wait()
123
124 return ini.out, ini.err
125 }
126
127 // TODO(bep) investigate if we can use sync.Cond for this.
128 func (ini *Init) wait() {
129 var counter time.Duration
130 for !ini.init.Done() {
131 counter += 10
132 if counter > 600000000 {
133 panic("BUG: timed out in lazy init")
134 }
135 time.Sleep(counter * time.Microsecond)
136 }
137 }
138
139 func (ini *Init) inProgress() bool {
140 return ini != nil && ini.init.InProgress()
141 }
142
143 func (ini *Init) shouldInitialize() bool {
144 return !(ini == nil || ini.init.Done() || ini.init.InProgress())
145 }
146
147 // Reset resets the current and all its dependencies.
148 func (ini *Init) Reset() {
149 mu := ini.init.ResetWithLock()
150 ini.err = nil
151 defer mu.Unlock()
152 for _, d := range ini.children {
153 d.Reset()
154 }
155 }
156
157 func (ini *Init) add(branch bool, initFn func() (any, error)) *Init {
158 ini.mu.Lock()
159 defer ini.mu.Unlock()
160
161 if branch {
162 return &Init{
163 f: initFn,
164 prev: ini,
165 }
166 }
167
168 ini.checkDone()
169 ini.children = append(ini.children, &Init{
170 f: initFn,
171 })
172
173 return ini
174 }
175
176 func (ini *Init) checkDone() {
177 if ini.init.Done() {
178 panic("init cannot be added to after it has run")
179 }
180 }
181
182 func (ini *Init) withTimeout(timeout time.Duration, f func(ctx context.Context) (any, error)) (any, error) {
183 ctx, cancel := context.WithTimeout(context.Background(), timeout)
184 defer cancel()
185 c := make(chan verr, 1)
186
187 go func() {
188 v, err := f(ctx)
189 select {
190 case <-ctx.Done():
191 return
192 default:
193 c <- verr{v: v, err: err}
194 }
195 }()
196
197 select {
198 case <-ctx.Done():
199 return nil, errors.New("timed out initializing value. You may have a circular loop in a shortcode, or your site may have resources that take longer to build than the `timeout` limit in your Hugo config file.")
200 case ve := <-c:
201 return ve.v, ve.err
202 }
203 }
204
205 type verr struct {
206 v any
207 err error
208 }