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 }