@@ -65,9 +65,7 @@ const (
6565 SDNRemoteArpMacAddress = "12-34-56-78-9a-bc"
6666
6767 // Command to fetch netadapter and pnp id
68- // TODO can we replace this (and things in endpoint_windows) with "golang.org/x/sys/windows"
69- // var adapterInfo windows.IpAdapterInfo
70- // var bufferSize uint32 = uint32(unsafe.Sizeof(adapterInfo))
68+ // TODO: can we replace this (and things in endpoint_windows) with other utils from "golang.org/x/sys/windows"?
7169 GetMacAddressVFPPnpIDMapping = "Get-NetAdapter | Select-Object MacAddress, PnpDeviceID| Format-Table -HideTableHeaders"
7270
7371 // Interval between successive checks for mellanox adapter's PriorityVLANTag value
@@ -248,9 +246,10 @@ func (p *execClient) ExecutePowershellCommandWithContext(ctx context.Context, co
248246}
249247
250248// SetSdnRemoteArpMacAddress sets the regkey for SDNRemoteArpMacAddress needed for multitenancy if hns is enabled
251- func SetSdnRemoteArpMacAddress () error {
249+ func SetSdnRemoteArpMacAddress (ctx context.Context ) error {
250+ log .Printf ("Setting SDNRemoteArpMacAddress regKey" )
252251 // open the registry key
253- k , err := registry .OpenKey (registry .LOCAL_MACHINE , " SYSTEM\\ CurrentControlSet\\ Services\\ hns\\ State" , registry .READ | registry .SET_VALUE )
252+ k , err := registry .OpenKey (registry .LOCAL_MACHINE , ` SYSTEM\CurrentControlSet\Services\hns\State` , registry .READ | registry .SET_VALUE )
254253 if err != nil {
255254 if errors .Is (err , registry .ErrNotExist ) {
256255 return nil
@@ -262,6 +261,8 @@ func SetSdnRemoteArpMacAddress() error {
262261 if err = setSDNRemoteARPMACAddress (k ); err != nil {
263262 return errors .Wrap (err , "could not set registry key" )
264263 }
264+ log .Printf ("SDNRemoteArpMacAddress regKey set successfully" )
265+ log .Printf ("Restarting HNS service" )
265266 // connect to the service manager
266267 m , err := mgr .Connect ()
267268 if err != nil {
@@ -274,44 +275,58 @@ func SetSdnRemoteArpMacAddress() error {
274275 return errors .Wrap (err , "could not access service" )
275276 }
276277 defer service .Close ()
277- return errors .Wrap (restartService (service ), "could not restart service" )
278+ if err := restartService (ctx , service ); err != nil {
279+ return errors .Wrap (err , "could not restart service" )
280+ }
281+ log .Printf ("HNS service restarted successfully" )
282+ return nil
278283}
279284
280285type key interface {
286+ GetStringValue (name string ) (val string , valtype uint32 , err error )
281287 SetStringValue (name , value string ) error
282288}
283289
284290func setSDNRemoteARPMACAddress (k key ) error {
285291 // Set the reg key
286292 // was "Set-ItemProperty -Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress -Value \"12-34-56-78-9a-bc\""
293+ if v , _ , _ := k .GetStringValue ("SDNRemoteArpMacAddress" ); v == SDNRemoteArpMacAddress {
294+ return nil // already set
295+ }
287296 if err := k .SetStringValue ("SDNRemoteArpMacAddress" , SDNRemoteArpMacAddress ); err != nil {
288297 return errors .Wrap (err , "could not set registry key" )
289298 }
290- log .Printf ("[Azure CNS] SDNRemoteArpMacAddress regKey set successfully. Restarting hns service." )
291299 return nil
292300}
293301
294- type serv interface {
295- Control (code svc.Cmd ) (svc.Status , error )
296- Query () (svc.Status , error )
297- Start (... string ) error
298- }
299-
300- func restartService (s serv ) error {
302+ func restartService (ctx context.Context , s * mgr.Service ) error {
301303 // Stop the service
302304 _ , err := s .Control (svc .Stop )
303305 if err != nil {
304306 return errors .Wrap (err , "could not stop service" )
305307 }
306308 // Wait for the service to stop
307- for status , err := s .Query (); status .State != svc .Stopped ; status , err = s .Query () {
309+ ticker := time .NewTicker (500 * time .Millisecond ) //nolint:gomnd // 500ms
310+ defer ticker .Stop ()
311+ for { // hacky cancellable do-while
312+ status , err := s .Query ()
308313 if err != nil {
309314 return errors .Wrap (err , "could not query service status" )
310315 }
311- time .Sleep (500 * time .Millisecond ) //nolint:gomnd // 500ms
316+ if status .State == svc .Stopped {
317+ break
318+ }
319+ select {
320+ case <- ctx .Done ():
321+ return errors .New ("context cancelled" )
322+ case <- ticker .C :
323+ }
312324 }
313325 // Start the service again
314- return errors .Wrap (s .Start (), "could not start service" )
326+ if err := s .Start (); err != nil {
327+ return errors .Wrap (err , "could not start service" )
328+ }
329+ return nil
315330}
316331
317332func HasMellanoxAdapter () bool {
@@ -384,7 +399,6 @@ func GetProcessNameByID(pidstr string) (string, error) {
384399 pidstr = strings .Trim (pidstr , "\r \n " )
385400 cmd := fmt .Sprintf ("Get-Process -Id %s|Format-List" , pidstr )
386401 p := NewExecClient (nil )
387- // TODO not riemovign this because it seems to only be called in test?
388402 out , err := p .ExecutePowershellCommand (cmd )
389403 if err != nil {
390404 log .Printf ("Process is not running. Output:%v, Error %v" , out , err )
0 commit comments