@@ -2,6 +2,7 @@ package utils
22
33import (
44 "fmt"
5+ "math"
56 "sort"
67)
78
@@ -17,121 +18,166 @@ type ServiceConfig struct {
1718 MaxMemory int
1819}
1920
20- type serviceLevel struct {
21- Priority int
22- Children []* ServiceConfig
21+ func minInt (a , b int ) int {
22+ if a < b {
23+ return a
24+ }
25+ return b
2326}
2427
25- func (s * serviceLevel ) balanceMemory (targetMemory int ) (int , error ) {
26- totalMemory := targetMemory
27- usedMemory := 0
28+ func maxInt (a , b int ) int {
29+ if a > b {
30+ return a
31+ }
32+ return b
33+ }
2834
29- for _ , child := range s .Children {
30- if child .MinMemory > totalMemory {
31- return usedMemory , fmt .Errorf ("not enough memory to satisfy minimum for service: %s" , child .Name )
32- }
33- child .AssignedMemory = child .MinMemory
34- totalMemory -= child .MinMemory
35- usedMemory += child .MinMemory
35+ var priorityWeight = map [int ]float64 {
36+ 1 : 4.0 ,
37+ 2 : 3.0 ,
38+ 3 : 2.0 ,
39+ }
40+
41+ func getWeight (priority int ) float64 {
42+ if weight , ok := priorityWeight [priority ]; ok {
43+ return weight
3644 }
45+ return 1.0
46+ }
3747
38- averageMemory := totalMemory / len (s .Children )
48+ func BalanceMemory (services []ServiceConfig , totalSystemMemory int ) (map [string ]* ServiceConfig , error ) {
49+ if totalSystemMemory <= SYSTEM_RESERVED_MEMORY {
50+ return nil , fmt .Errorf ("total system memory (%dMB) is not greater than reserved memory (%dMB)" , totalSystemMemory , SYSTEM_RESERVED_MEMORY )
51+ }
52+ availableMemory := totalSystemMemory - SYSTEM_RESERVED_MEMORY
3953
40- for _ , child := range s . Children {
41- assignableMemory := min ( averageMemory , child . MaxMemory - child . AssignedMemory )
42- child . AssignedMemory += assignableMemory
43- totalMemory -= assignableMemory
44- usedMemory += assignableMemory
54+ workingServices := make ([] * ServiceConfig , len ( services ))
55+ for i := range services {
56+ s := services [ i ]
57+ workingServices [ i ] = & s
58+ workingServices [ i ]. AssignedMemory = 0
4559 }
4660
47- noLimitChildren := [] * ServiceConfig {}
48- for _ , child := range s . Children {
49- if child . MaxMemory == 0 {
50- noLimitChildren = append ( noLimitChildren , child )
61+ totalMinMemoryNeeded := 0
62+ for _ , s := range workingServices {
63+ if s . MinMemory < 0 {
64+ return nil , fmt . Errorf ( "service %s has negative MinMemory: %dMB" , s . Name , s . MinMemory )
5165 }
52- }
53- if len (noLimitChildren ) > 0 {
54- memoryPerChild := totalMemory / len (noLimitChildren )
55- for _ , child := range noLimitChildren {
56- child .AssignedMemory += memoryPerChild
57- totalMemory -= memoryPerChild
58- usedMemory += memoryPerChild
66+ if s .MaxMemory != 0 && s .MinMemory > s .MaxMemory {
67+ return nil , fmt .Errorf ("service %s has MinMemory (%dMB) greater than MaxMemory (%dMB)" , s .Name , s .MinMemory , s .MaxMemory )
5968 }
69+ totalMinMemoryNeeded += s .MinMemory
6070 }
6171
62- return usedMemory , nil
63- }
72+ if totalMinMemoryNeeded > availableMemory {
73+ return nil , fmt .Errorf ("insufficient memory: Available %dMB < Total Minimum Required %dMB. (System requires %dMB)" ,
74+ availableMemory , totalMinMemoryNeeded , totalMinMemoryNeeded + SYSTEM_RESERVED_MEMORY )
75+ }
6476
65- func ( s * serviceLevel ) getMinimum () int {
66- totalMinMemory := 0
67- for _ , child := range s . Children {
68- totalMinMemory += child .MinMemory
77+ memoryUsed := 0
78+ for _ , s := range workingServices {
79+ s . AssignedMemory = s . MinMemory
80+ memoryUsed += s .MinMemory
6981 }
82+ remainingMemory := availableMemory - memoryUsed
83+
84+ for remainingMemory > 0 {
85+ distributableInThisPass := remainingMemory
86+ candidates := []* ServiceConfig {}
87+ totalWeight := 0.0
88+
89+ for _ , s := range workingServices {
90+ canTakeMore := (s .MaxMemory == 0 ) || (s .AssignedMemory < s .MaxMemory )
91+ if canTakeMore {
92+ candidates = append (candidates , s )
93+ totalWeight += getWeight (s .Priority )
94+ }
95+ }
7096
71- return totalMinMemory
72- }
97+ if len (candidates ) == 0 || totalWeight <= 0 {
98+ break
99+ }
73100
74- func balanceMemoryAcrossTrees (trees []* serviceLevel , totalMemory int ) error {
75- totalMinMemory := 0
76- for _ , tree := range trees {
77- totalMinMemory += tree .getMinimum ()
78- }
79- if totalMemory < totalMinMemory {
80- return fmt .Errorf ("your system does not have the minimum required memory: %dMB" , totalMinMemory + SYSTEM_RESERVED_MEMORY )
81- }
101+ memoryAssignedInPass := 0
102+ sort .SliceStable (candidates , func (i , j int ) bool {
103+ return candidates [i ].Priority < candidates [j ].Priority
104+ })
82105
83- sort .Slice (trees , func (i , j int ) bool {
84- return trees [i ].Priority < trees [j ].Priority
85- })
106+ tempAssignments := make (map [string ]int )
107+ for _ , s := range candidates {
108+ weight := getWeight (s .Priority )
109+ proportionalShare := float64 (distributableInThisPass ) * (weight / totalWeight )
86110
87- targetPercentages := []float64 {0.7 , 0.2 , 0.1 }
111+ remainingCapacity := math .MaxInt
112+ if s .MaxMemory != 0 {
113+ remainingCapacity = s .MaxMemory - s .AssignedMemory
114+ }
88115
89- for i , tree := range trees {
90- targetMemory := int (float64 (totalMemory - totalMinMemory ) * targetPercentages [i ])
91- if targetMemory > totalMemory {
92- targetMemory = totalMemory
93- }
116+ assignAmount := int (math .Floor (proportionalShare ))
117+ assignAmount = minInt (assignAmount , remainingCapacity )
118+ assignAmount = maxInt (0 , assignAmount )
94119
95- _ , err := tree .balanceMemory (tree .getMinimum () + targetMemory )
96- if err != nil {
97- fmt .Printf ("err: %v\n " , err )
120+ tempAssignments [s .Name ] = assignAmount
121+ memoryAssignedInPass += assignAmount
98122 }
99- }
100123
101- return nil
102- }
124+ for _ , s := range candidates {
125+ s .AssignedMemory += tempAssignments [s .Name ]
126+ }
103127
104- func createServiceLevels (services []ServiceConfig ) []* serviceLevel {
105- levelTrees := make (map [int ]* serviceLevel )
128+ remainingMemory -= memoryAssignedInPass
106129
107- for _ , sv := range services {
108- sv := sv
109- if _ , ok := levelTrees [sv .Priority ]; ! ok {
110- levelTrees [sv .Priority ] = & serviceLevel {Priority : sv .Priority , Children : []* ServiceConfig {}}
130+ if memoryAssignedInPass == 0 && remainingMemory > 0 {
131+ break
111132 }
112- levelTrees [sv .Priority ].Children = append (levelTrees [sv .Priority ].Children , & sv )
113133 }
114134
115- trees := []* serviceLevel {}
116- for _ , tree := range levelTrees {
117- trees = append (trees , tree )
135+ if remainingMemory > 0 {
136+ sort .SliceStable (workingServices , func (i , j int ) bool {
137+ if workingServices [i ].Priority != workingServices [j ].Priority {
138+ return workingServices [i ].Priority < workingServices [j ].Priority
139+ }
140+ if workingServices [i ].MaxMemory == 0 && workingServices [j ].MaxMemory != 0 {
141+ return true
142+ }
143+ if workingServices [i ].MaxMemory != 0 && workingServices [j ].MaxMemory == 0 {
144+ return false
145+ }
146+ return workingServices [i ].Name < workingServices [j ].Name
147+ })
148+
149+ for remainingMemory > 0 {
150+ memoryDistributedInSweep := false
151+ for _ , s := range workingServices {
152+ if remainingMemory == 0 {
153+ break
154+ }
155+ canTakeMore := (s .MaxMemory == 0 ) || (s .AssignedMemory < s .MaxMemory )
156+ if canTakeMore {
157+ s .AssignedMemory ++
158+ remainingMemory --
159+ memoryDistributedInSweep = true
160+ }
161+ }
162+ if ! memoryDistributedInSweep && remainingMemory > 0 {
163+ fmt .Printf ("WARNING: Unable to distribute final %dMB in sweep phase. Memory might be underutilized.\n " , remainingMemory )
164+ break
165+ }
166+ }
118167 }
119168
120- return trees
121- }
122-
123- func BalanceMemory (services []ServiceConfig , totalMemory int ) (map [string ]* ServiceConfig , error ) {
124- trees := createServiceLevels (services )
125- err := balanceMemoryAcrossTrees (trees , totalMemory )
126- if err != nil {
127- return nil , fmt .Errorf ("error distributing memory: %v" , err )
169+ finalAssignedMemory := 0
170+ serviceMap := make (map [string ]* ServiceConfig )
171+ for _ , s := range workingServices {
172+ serviceMap [s .Name ] = s
173+ finalAssignedMemory += s .AssignedMemory
128174 }
129175
130- serviceMap := make ( map [ string ] * ServiceConfig )
131- for _ , tree := range trees {
132- for _ , service := range tree . Children {
133- serviceMap [ service . Name ] = service
134- }
176+ if finalAssignedMemory != availableMemory {
177+ fmt . Printf ( "WARNING: Final memory validation failed. Expected %dMB, Assigned %dMB. Discrepancy: %dMB \n " ,
178+ availableMemory , finalAssignedMemory , availableMemory - finalAssignedMemory )
179+ } else {
180+ fmt . Printf ( "Memory distribution successful. Total Assigned: %dMB (Available: %dMB) \n " , finalAssignedMemory , availableMemory )
135181 }
136182
137183 return serviceMap , nil
0 commit comments