Skip to content

Commit

Permalink
added multiple radium server support
Browse files Browse the repository at this point in the history
  • Loading branch information
spy16 committed Sep 23, 2018
1 parent 4833bec commit 5c06ab5
Show file tree
Hide file tree
Showing 5 changed files with 38 additions and 8 deletions.
8 changes: 4 additions & 4 deletions cmd/radium/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,8 @@ import (

func initConfig(cfg *config, rootCmd *cobra.Command) {
viper.SetDefault("sources", "cheatsh,learnxiny,tldr")
viper.SetDefault("radiumurl", "")
viper.RegisterAlias("radiumurl", "radium_url")
viper.SetDefault("radiumservers", []string{})
viper.RegisterAlias("radiumservers", "radium_servers")

viper.SetConfigName("radium")
viper.AddConfigPath("./")
Expand All @@ -30,6 +30,6 @@ func initConfig(cfg *config, rootCmd *cobra.Command) {
// config struct is used to store CLI configurations. configuration
// values are read into this struct using viper
type config struct {
Sources []string `json:"sources"`
RadiumURL string `json:"radium_url"`
Sources []string `json:"sources"`
RadiumServers []string `json:"radium_servers"`
}
9 changes: 9 additions & 0 deletions cmd/radium/query.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"fmt"
"os"
"strings"
"time"

"github.com/shivylp/radium"
"github.com/spf13/cobra"
Expand All @@ -20,7 +21,9 @@ func newQueryCmd(cfg *config) *cobra.Command {
}

var attribs []string
var deadline time.Duration
cmd.Flags().StringSliceVarP(&attribs, "attr", "a", []string{}, "Attributes to narrow the search scope")
cmd.Flags().DurationVarP(&deadline, "timeout", "t", 10*time.Second, "Timeout to use across all sources")

cmd.Run = func(_ *cobra.Command, args []string) {
query := radium.Query{}
Expand All @@ -45,6 +48,12 @@ func newQueryCmd(cfg *config) *cobra.Command {
}

ctx := context.Background()
var cancel func()
if deadline.Seconds() > 0 {
ctx, cancel = context.WithTimeout(ctx, deadline)
defer cancel()
}

ins := getNewRadiumInstance(*cfg)
rs, err := ins.Search(ctx, query, strategy)
if err != nil {
Expand Down
10 changes: 9 additions & 1 deletion cmd/radium/setup.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package main

import (
"fmt"
"strings"

"github.com/shivylp/radium"
Expand All @@ -26,7 +27,14 @@ func getNewRadiumInstance(cfg config) *radium.Instance {
case "duckduckgo", "ddg":
ins.RegisterSource("duckduckgo", duckduckgo.New())
case "radium", "rad":
ins.RegisterSource("radium", sources.NewRadium(cfg.RadiumURL))
if len(cfg.RadiumServers) == 0 {
ins.Logger.Warnf("radium source is enabled, but no radium_server specified")
} else {
for id, url := range cfg.RadiumServers {
name := fmt.Sprintf("radium-%d", id)
ins.RegisterSource(name, sources.NewRadium(url))
}
}
default:
ins.Fatalf("unknown source type: %s", src)
}
Expand Down
8 changes: 7 additions & 1 deletion radium.yaml
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
# enabled list of sources. Sources will be executed one by one
# in the order they are specified when using '1st' strategy.
# Order is ignored when using 'concurrent' strategy.
sources:
- cheat.sh
- tldr
Expand All @@ -6,4 +9,7 @@ sources:
- duckduckgo
- radium

radium_url: http://localhost:8080
# list of radium servers to be configured. used only when
# the 'sources' configuration contains 'radium'
radium_servers:
- http://localhost:8080
11 changes: 9 additions & 2 deletions strategy.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,17 +30,24 @@ func (con Concurrent) Execute(ctx context.Context, query Query, sources []Regist
wg := &sync.WaitGroup{}

for _, source := range sources {
wg.Add(1)
select {
case <-ctx.Done():
con.Infof("received cancel signal. stopping")
break
default:
}

wg.Add(1)
go func(wg *sync.WaitGroup, src RegisteredSource, rs *safeResults) {
defer wg.Done()

srcResults, err := src.Search(ctx, query)
if err != nil {
con.Warnf("source '%s' failed: %s", src.Name, err)
return
}

rs.extend(srcResults, src.Name, con.Logger)
wg.Done()
}(wg, source, results)
}

Expand Down

0 comments on commit 5c06ab5

Please sign in to comment.