diff --git a/cmd/microcloud/main_init.go b/cmd/microcloud/main_init.go index 323381a30..80b8af2db 100644 --- a/cmd/microcloud/main_init.go +++ b/cmd/microcloud/main_init.go @@ -192,14 +192,13 @@ func lookupPeers(s *service.Handler, autoSetup bool, iface *net.Interface, subne }() } - var timeAfter <-chan time.Time + timeoutDuration := time.Minute if autoSetup { - timeAfter = time.After(5 * time.Second) + timeoutDuration = 5 * time.Second } - if len(expectedSystems) > 0 { - timeAfter = time.After(1 * time.Minute) - } + ctx, cancel := context.WithTimeout(context.Background(), timeoutDuration) + defer cancel() expectedSystemsMap := make(map[string]bool, len(expectedSystems)) for _, system := range expectedSystems { @@ -211,7 +210,7 @@ func lookupPeers(s *service.Handler, autoSetup bool, iface *net.Interface, subne done := false for !done { select { - case <-timeAfter: + case <-ctx.Done(): done = true case err := <-selectionCh: if err != nil { @@ -227,7 +226,7 @@ func lookupPeers(s *service.Handler, autoSetup bool, iface *net.Interface, subne break } - peers, err := mdns.LookupPeers(context.Background(), iface, mdns.Version, s.Name) + peers, err := mdns.LookupPeers(ctx, iface, mdns.Version, s.Name) if err != nil { return err }