├── .gitignore ├── AUTHORS ├── LICENSE ├── README.md ├── mock_client.go ├── solrclient.go ├── solrclient_test.go ├── solrjob.go ├── solrquery.go ├── solrquery_test.go ├── solrresponse.go ├── solrresponse_test.go ├── workerpool.go └── workerpool_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled Object files, Static and Dynamic libs (Shared Objects) 2 | *.o 3 | *.a 4 | *.so 5 | 6 | # Folders 7 | _obj 8 | _test 9 | 10 | # Architecture specific extensions/prefixes 11 | *.[568vq] 12 | [568vq].out 13 | 14 | *.cgo1.go 15 | *.cgo2.c 16 | _cgo_defun.c 17 | _cgo_gotypes.go 18 | _cgo_export.* 19 | 20 | _testmain.go 21 | 22 | *.exe 23 | *.test 24 | *.prof 25 | -------------------------------------------------------------------------------- /AUTHORS: -------------------------------------------------------------------------------- 1 | Srdjan Marinovic, srdjan.marinovic@gmail.com, @a-little-srdjan 2 | Ryan Day, ryan@ryanday.net, @rday -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2016 wirelessregistry 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # gora 2 | 3 | A simple, bare-bones solr client for Go. 4 | 5 | This is a driver in-the-making, and should be approached as such. We appreciate any feedback via issues. Please be harsh and do break our code. 6 | 7 | Key features: 8 | * Uses HTTP connections. 9 | * Support for facets and standard queries. 10 | * Support for multiple hosts (SolrCloud) with recovery. 11 | 12 | ### Installation 13 | 14 | ``` 15 | go get github.com/wirelessregistry/gora 16 | ``` 17 | 18 | ### Usage 19 | 20 | Gora consists of two layers: 21 | * A low-level solr client. (see solrclient.go) 22 | * A connection-management pool. (see workerpool.go) 23 | 24 | __Low-level client__ 25 | 26 | This is synchronous client, that simply constructs an appropriate JSON post to a given node+core combo. Its input is a SolrJob interface, and it deserializes the resulting JSON as a SolrResponse struct. 27 | 28 | The SolrJob interface provides a Bytes() function, which the client uses to obtain the raw query to send to the Solr server. A Handler() function is also provided, giving the job control over how it is processed by Solr. 29 | 30 | See solrclient_test.go for usage examples. 31 | 32 | __Connection Pool__ 33 | 34 | The key inconvenience when using the aforementioned client is concurrent processing and connection management. In detail: 35 | 36 | One can launch multiple goroutines (e.g. in the master-slave pattern) to execute queries concurrently. This approach works well when a process does not launch excessive numbers of goroutines. When this does not hold, the connection pool can be launched with a fixed number of running goroutines. In this case, a process submits a job to the pool and awaits the query completion. 37 | 38 | A pool can be started with a set of solr hosts (e.g. SolrCloud). In this case, equal number of goroutines will be dedicated to each host. Note that the sharding strategy is not taken into account when assigning jobs to routines. If a host becomes unavailable, the corresponding routines will take themselves offline and wait until the host is again available before taking on new jobs. 39 | 40 | -------------------------------------------------------------------------------- /mock_client.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | import "strconv" 4 | 5 | type MockSolrClient struct { 6 | Timeout bool 7 | ConnectCount int 8 | CloseCh chan struct{} 9 | Response *SolrResponse 10 | } 11 | 12 | func (c *MockSolrClient) Execute(s SolrJob) (*SolrResponse, bool) { 13 | c.Response.Status, _ = strconv.Atoi(string(s.Bytes())) 14 | 15 | // Client will keep iterating through the result set until 16 | // no additional documents are returned. 17 | if c.Response != nil && c.Response.Response != nil { 18 | if s.GetStart() > len(c.Response.Response.Docs) { 19 | res := &SolrResponse{ 20 | Response: &DocumentCollection{}, 21 | Status: c.Response.Status, 22 | } 23 | 24 | return res, c.Timeout 25 | } 26 | } 27 | 28 | return c.Response, c.Timeout 29 | } 30 | 31 | func (c *MockSolrClient) TestConnection() bool { 32 | c.ConnectCount += 1 33 | 34 | if c.ConnectCount > 3 { 35 | c.CloseCh <- struct{}{} 36 | } 37 | 38 | return c.Timeout 39 | } 40 | 41 | func MockClientConstructor(hostUrl string, core string, ch chan struct{}) SolrClient { 42 | return &MockSolrClient{ 43 | Timeout: false, 44 | CloseCh: ch, 45 | Response: &SolrResponse{}, 46 | } 47 | } 48 | 49 | type MockSolrJob struct { 50 | payload []byte 51 | rCh chan *SolrResponse 52 | } 53 | 54 | func NewMockSolrJob(payload []byte) *MockSolrJob { 55 | return &MockSolrJob{ 56 | payload: payload, 57 | rCh: make(chan *SolrResponse, 1), 58 | } 59 | } 60 | 61 | func (j *MockSolrJob) Handler() string { 62 | return "mochHandler" 63 | } 64 | 65 | func (j *MockSolrJob) Bytes() []byte { 66 | return j.payload 67 | } 68 | 69 | func (j *MockSolrJob) ResultCh() chan *SolrResponse { 70 | return j.rCh 71 | } 72 | 73 | func (j *MockSolrJob) Wait() *SolrResponse { 74 | return <-j.ResultCh() 75 | } 76 | 77 | func (j *MockSolrJob) GetRows() int { 78 | return 0 79 | } 80 | 81 | func (j *MockSolrJob) GetStart() int { 82 | return 0 83 | } 84 | -------------------------------------------------------------------------------- /solrclient.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "io/ioutil" 7 | "net/http" 8 | "net/url" 9 | 10 | "github.com/wirelessregistry/glog" 11 | ) 12 | 13 | // SolrClient is the interface that workers in the workerpool use 14 | // to get SolrResponses for SolrJobs. 15 | // 16 | // Execute(SolrJob) should always return a SolrRespose. If there was 17 | // an error executing the SolrJob, the error should be put in the 18 | // SolrResponse, and the retry flag should be set depending on whether 19 | // this error is recoverable or not. 20 | // 21 | // TestConnection() will be called by the worker if an error has 22 | // previously occurred. The worker will not accept any new jobs 23 | // until the SolrClient has a vsalid connection. 24 | type SolrClient interface { 25 | Execute(SolrJob) (*SolrResponse, bool) 26 | TestConnection() bool 27 | } 28 | 29 | type HttpSolrClient struct { 30 | // Host specifies the URL of the Solr Server 31 | Host string 32 | 33 | // Core specifies the Solr core to work with 34 | Core string 35 | 36 | username string 37 | password string 38 | 39 | client *http.Client 40 | } 41 | 42 | // NewHttpSolrClient creates a SolrClient with an http.Client connection 43 | func NewHttpSolrClient(host, core string) SolrClient { 44 | transport := http.Transport{ 45 | MaxIdleConnsPerHost: 2, 46 | } 47 | 48 | httpClient := http.Client{ 49 | Transport: &transport, 50 | } 51 | 52 | client := HttpSolrClient{ 53 | Host: host, 54 | Core: core, 55 | client: &httpClient, 56 | } 57 | 58 | return &client 59 | } 60 | 61 | // NewHttpSolrClient creates a SolrClient with an http.Client connection and uses basic authentication. 62 | func NewHttpSolrClientWithAuth(host, core, username, password string) SolrClient { 63 | transport := http.Transport{ 64 | MaxIdleConnsPerHost: 2, 65 | } 66 | 67 | httpClient := http.Client{ 68 | Transport: &transport, 69 | } 70 | 71 | client := HttpSolrClient{ 72 | Host: host, 73 | Core: core, 74 | client: &httpClient, 75 | username: username, 76 | password: password, 77 | } 78 | 79 | return &client 80 | } 81 | 82 | func (c *HttpSolrClient) useAuth() bool { 83 | return (len(c.username) > 0 && len(c.password) > 0) 84 | } 85 | 86 | // TestConnection will issue an empty query to the Solr server. 87 | // As long as we don't get an error, we know that the Solr server 88 | // received the query, and that this connection is valid. 89 | func (c *HttpSolrClient) TestConnection() bool { 90 | _, err := c.execQuery("", []byte("")) 91 | 92 | if err != nil && glog.V(2) { 93 | glog.Infof("HttpSolrClient.TestConnection() for %v failed. %v.", c.Host, err) 94 | } 95 | 96 | return err == nil 97 | } 98 | 99 | // Execute will send the given job to the Solr server and wait for 100 | // a response. If an error is received, the retry value will be determined 101 | // and the error will be placed in an empty SolrResponse. 102 | func (c *HttpSolrClient) Execute(job SolrJob) (*SolrResponse, bool) { 103 | handler := job.Handler() 104 | jobBytes := job.Bytes() 105 | 106 | emptyResponse := &SolrResponse{} 107 | byteResponse, err := c.execQuery(handler, jobBytes) 108 | if err != nil { 109 | glog.Warningf("HttpSolrClient.execQuery() failed. %v.", err) 110 | 111 | emptyResponse.Error = err 112 | return emptyResponse, c.temporaryError(err) 113 | } 114 | 115 | solrResponse, err := SolrResponseFromHTTPResponse(byteResponse) 116 | if err != nil { 117 | glog.Errorf("HttpSolrClient.SolrResponseFromHTTPResponse() failed. %v.", err) 118 | glog.Errorf("Found %v", string(byteResponse)) 119 | 120 | emptyResponse.Error = err 121 | return emptyResponse, false 122 | } 123 | 124 | return solrResponse, false 125 | } 126 | 127 | // temporaryError tries to determine whether this error is recoverable. 128 | // Most errors will be type *url.Error, and we can ask that error 129 | // whether it is temporary, or timeout related. 130 | func (c *HttpSolrClient) temporaryError(err error) bool { 131 | if err == nil { 132 | return false 133 | } 134 | 135 | // Could have an internal buffer problem, we should try again 136 | if err == bytes.ErrTooLarge { 137 | return true 138 | } 139 | 140 | // If we get a totally unknown error, send it back to the caller 141 | urlError, ok := err.(*url.Error) 142 | if !ok { 143 | return false 144 | } 145 | 146 | return urlError.Temporary() || urlError.Timeout() 147 | } 148 | 149 | // execQuery creates the full URL and posts an array of bytes to that url. 150 | func (c *HttpSolrClient) execQuery(handler string, json []byte) ([]byte, error) { 151 | url := fmt.Sprintf("%s/solr/%s/%s", c.Host, c.Core, handler) 152 | 153 | req, err := http.NewRequest("POST", url, bytes.NewReader(json)) 154 | if err != nil { 155 | return nil, err 156 | } 157 | req.Header.Set("Content-Type", "application/json") 158 | 159 | if c.useAuth() { 160 | req.SetBasicAuth(c.username, c.password) 161 | } 162 | 163 | r, err := c.client.Do(req) 164 | if err != nil { 165 | return nil, err 166 | } 167 | 168 | defer r.Body.Close() 169 | 170 | defer r.Body.Close() 171 | 172 | // read the response and check 173 | body, err := ioutil.ReadAll(r.Body) 174 | if err != nil { 175 | return nil, err 176 | } 177 | 178 | return body, nil 179 | } 180 | -------------------------------------------------------------------------------- /solrclient_test.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | import ( 4 | "bytes" 5 | "io" 6 | "io/ioutil" 7 | "net/http" 8 | "net/http/httptest" 9 | "net/url" 10 | "path" 11 | "strings" 12 | "testing" 13 | "time" 14 | ) 15 | 16 | // RewriteTransport is an http.RoundTripper that rewrites requests 17 | // using the provided URL's Scheme and Host, and its Path as a prefix. 18 | // The Opaque field is untouched. 19 | // If Transport is nil, http.DefaultTransport is used 20 | // http://stackoverflow.com/a/27894872/61980 21 | type RewriteTransport struct { 22 | // Dial net.Dial 23 | ResponseHeaderTimeout time.Duration 24 | Transport http.RoundTripper 25 | URL *url.URL 26 | } 27 | 28 | func (t RewriteTransport) RoundTrip(req *http.Request) (*http.Response, error) { 29 | // note that url.URL.ResolveReference doesn't work here 30 | // since t.u is an absolute url 31 | req.URL.Scheme = t.URL.Scheme 32 | req.URL.Host = t.URL.Host 33 | req.URL.Path = path.Join(t.URL.Path, req.URL.Path) 34 | rt := t.Transport 35 | if rt == nil { 36 | rt = http.DefaultTransport 37 | } 38 | 39 | return rt.RoundTrip(req) 40 | } 41 | 42 | func createTestServer(testData io.Reader, expectedMethod string) (*httptest.Server, *HttpSolrClient) { 43 | testBuf, _ := ioutil.ReadAll(testData) 44 | testData.Read(testBuf) 45 | 46 | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 47 | buf := bytes.NewBuffer(testBuf) 48 | if !strings.HasSuffix(r.URL.String(), expectedMethod) { 49 | http.NotFound(w, r) 50 | return 51 | } 52 | io.Copy(w, buf) 53 | }) 54 | 55 | server := httptest.NewServer(handler) 56 | testUrl, _ := url.Parse(server.URL) 57 | 58 | client := http.Client{ 59 | Transport: RewriteTransport{ 60 | URL: testUrl, 61 | }, 62 | } 63 | 64 | solrClient := NewHttpSolrClient(server.URL, "core").(*HttpSolrClient) 65 | solrClient.client = &client 66 | 67 | // Return our test server and a client which points to it 68 | return server, solrClient 69 | } 70 | 71 | func TestQuery(t *testing.T) { 72 | expected := bytes.NewBufferString(`{ 73 | "responseHeader": { 74 | "status": 0, 75 | "QTime": 72, 76 | "params": { 77 | "q": "*:*", 78 | "indent": "true", 79 | "wt": "json", 80 | "_": "1456851532479" 81 | } 82 | }, 83 | "response": { 84 | "numFound": 21, 85 | "start": 0, 86 | "docs": [ 87 | { 88 | "id": "My Id", 89 | "greeting": "你好", 90 | "list": [ 91 | "Jedná se o delší položka" 92 | ], 93 | "_version_": 1525161299814645800 94 | } 95 | ] 96 | } 97 | }`) 98 | 99 | server, client := createTestServer(expected, "/select") 100 | defer server.Close() 101 | 102 | solrQuery := NewSolrQuery("*:*", 0, 100, nil, nil, nil, "/select") 103 | resp, retry := client.Execute(solrQuery) 104 | if retry { 105 | t.Fatalf("Should not have to retry a valid job") 106 | } 107 | 108 | if resp != nil && resp.Response.NumFound != 21 { 109 | t.Errorf("Expected 21 documents, found %d", resp.Response.NumFound) 110 | } 111 | } 112 | 113 | func TestErrorQuery(t *testing.T) { 114 | expected := bytes.NewBufferString(`{ 115 | "responseHeader": { 116 | "status": 400, 117 | "QTime": 1 118 | }, 119 | "error": { 120 | "msg": "ERROR: [doc=change.me] unknown field 'title'", 121 | "code": 400 122 | } 123 | }`) 124 | 125 | server, client := createTestServer(expected, "/select") 126 | defer server.Close() 127 | 128 | solrQuery := NewSolrQuery("*:*", 0, 100, nil, nil, nil, "/select") 129 | resp, retry := client.Execute(solrQuery) 130 | if retry { 131 | t.Fatalf("Should not have to retry a valid job") 132 | } 133 | 134 | if resp != nil && resp.Status != 400 { 135 | t.Errorf("Expected 400 status, found %d", resp.Status) 136 | } 137 | 138 | if resp != nil && resp.Error.Error() != "ERROR: [doc=change.me] unknown field 'title'" { 139 | t.Errorf("Expected error message status, found %s", resp.Error) 140 | } 141 | } 142 | 143 | func TestTestConnection(t *testing.T) { 144 | expected := bytes.NewBufferString(`{}`) 145 | server, client := createTestServer(expected, "/update") 146 | defer server.Close() 147 | 148 | working := client.TestConnection() 149 | if !working { 150 | t.Error("Connection should be working") 151 | } 152 | 153 | handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) 154 | server = httptest.NewServer(handler) 155 | testUrl, _ := url.Parse("http://127.0.0.2:1/") 156 | 157 | httpClient := http.Client{ 158 | Transport: RewriteTransport{ 159 | URL: testUrl, 160 | }, 161 | } 162 | 163 | solrClient := NewHttpSolrClient(server.URL, "core").(*HttpSolrClient) 164 | solrClient.client = &httpClient 165 | working = solrClient.TestConnection() 166 | if working { 167 | t.Error("Connection should not be working") 168 | } 169 | 170 | } 171 | 172 | func TestBadDecode(t *testing.T) { 173 | expected := bytes.NewBufferString(`{ 174 | "responseHeader": { 175 | "status": 0, 176 | "QTime": 72, 177 | }`) 178 | 179 | server, client := createTestServer(expected, "/update") 180 | defer server.Close() 181 | 182 | solrQuery := NewSolrQuery("*:*", 0, 100, nil, nil, nil, "/update") 183 | _, retry := client.Execute(solrQuery) 184 | if retry { 185 | t.Error("We should not retry this job") 186 | } 187 | } 188 | -------------------------------------------------------------------------------- /solrjob.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | // SolrJob is the interface that a SolrClient needs in order 4 | // to get enough information to connnect to a Solr server. 5 | type SolrJob interface { 6 | // Handler returns the method used to run the job. (select, update) 7 | Handler() string 8 | 9 | // Bytes returns the array of bytes representing the JSON query 10 | Bytes() []byte 11 | 12 | // ResultCh return the channel that the job's response will be sent to 13 | ResultCh() chan *SolrResponse 14 | 15 | // Wait is a convinience method that allows a one line function 16 | // to wait for a SolrResponse 17 | Wait() *SolrResponse 18 | 19 | GetRows() int 20 | GetStart() int 21 | } 22 | -------------------------------------------------------------------------------- /solrquery.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | import ( 4 | "bytes" 5 | "encoding/json" 6 | "fmt" 7 | "strconv" 8 | "strings" 9 | 10 | "github.com/wirelessregistry/glog" 11 | ) 12 | 13 | // SolrQuery represents a SolrJob that can be submitted to a pool. 14 | // It is a bog standard query representation. 15 | type SolrQuery struct { 16 | Rows int 17 | Start int 18 | Query string 19 | Facet *string // this has to be a raw json facet query! 20 | Filter *string // this has to be a raw json filter query! 21 | Sort *string // field order 22 | Params map[string]interface{} 23 | handler string 24 | resultCh chan *SolrResponse 25 | } 26 | 27 | func NewSolrQuery(q string, s, r int, filter, facet, sort *string, handler string) *SolrQuery { 28 | params := make(map[string]interface{}) 29 | params["wt"] = "json" 30 | 31 | return &SolrQuery{ 32 | Query: q, 33 | Start: s, 34 | Rows: r, 35 | Facet: facet, 36 | Filter: filter, 37 | Sort: sort, 38 | Params: params, 39 | handler: handler, 40 | resultCh: make(chan *SolrResponse, 1), 41 | } 42 | } 43 | 44 | func NewSolrSpatialQuery(q, stype, sfield string, lat, lon, d float64, s, r int, filter, facet, sort *string, handler string) *SolrQuery { 45 | params := make(map[string]interface{}) 46 | params["wt"] = "json" 47 | params["fq"] = fmt.Sprintf("{!%s sfield=%s}", stype, sfield) 48 | params["pt"] = fmt.Sprintf("%f,%f", lat, lon) 49 | params["d"] = fmt.Sprintf("%f", d) 50 | 51 | return &SolrQuery{ 52 | Query: q, 53 | Start: s, 54 | Rows: r, 55 | Facet: facet, 56 | Filter: filter, 57 | Sort: sort, 58 | Params: params, 59 | handler: handler, 60 | resultCh: make(chan *SolrResponse, 1), 61 | } 62 | } 63 | 64 | func escape(s string) string { 65 | return "\"" + s + "\"" 66 | } 67 | 68 | func (q *SolrQuery) Handler() string { 69 | return q.handler 70 | } 71 | 72 | func (q *SolrQuery) ResultCh() chan *SolrResponse { 73 | return q.resultCh 74 | } 75 | 76 | func (q *SolrQuery) Wait() *SolrResponse { 77 | return <-q.ResultCh() 78 | } 79 | 80 | func (q *SolrQuery) GetRows() int { 81 | return q.Rows 82 | } 83 | 84 | func (q *SolrQuery) GetStart() int { 85 | return q.Start 86 | } 87 | 88 | // Bytes will manually construct the JSON query that will be sent 89 | // to the Solr server. We do this to embed additional JSON fields. 90 | // This function executes a tad faster than using the json.Marshal function. 91 | func (q *SolrQuery) Bytes() []byte { 92 | query := make(map[string]interface{}) 93 | 94 | query["query"] = q.Query 95 | 96 | if q.Sort != nil { 97 | query["sort"] = *q.Sort 98 | } 99 | 100 | if q.Filter != nil { 101 | query["filter"] = *q.Filter 102 | } 103 | 104 | if q.Facet != nil { 105 | query["facet"] = *q.Facet 106 | } 107 | 108 | q.Params["start"] = q.Start 109 | q.Params["rows"] = q.Rows 110 | 111 | query["params"] = q.Params 112 | 113 | b, err := json.Marshal(query) 114 | if err != nil { 115 | glog.Error(err) 116 | } 117 | 118 | return b 119 | } 120 | 121 | // SolrUpdateQuery represents a query that will update or create a new Solr document 122 | type SolrUpdateQuery struct { 123 | Documents map[string]interface{} 124 | handler string 125 | resultCh chan *SolrResponse 126 | } 127 | 128 | func NewSolrUpdateQuery(document map[string]interface{}) *SolrUpdateQuery { 129 | return &SolrUpdateQuery{ 130 | Documents: document, 131 | handler: "update", 132 | resultCh: make(chan *SolrResponse, 1), 133 | } 134 | } 135 | 136 | func (q *SolrUpdateQuery) Handler() string { 137 | return q.handler 138 | } 139 | 140 | func (q *SolrUpdateQuery) ResultCh() chan *SolrResponse { 141 | return q.resultCh 142 | } 143 | 144 | func (q *SolrUpdateQuery) Wait() *SolrResponse { 145 | return <-q.ResultCh() 146 | } 147 | 148 | func (q *SolrUpdateQuery) GetRows() int { 149 | return 0 150 | } 151 | 152 | func (q *SolrUpdateQuery) GetStart() int { 153 | return 0 154 | } 155 | 156 | func (q *SolrUpdateQuery) Bytes() []byte { 157 | b, _ := json.Marshal(q.Documents) 158 | buffer := bytes.NewBufferString(fmt.Sprintf("{\"add\":{\"doc\":%s}, \"commit\": {}}", b)) 159 | 160 | return buffer.Bytes() 161 | } 162 | 163 | // SolrBatchUpdateQuery represents a query that will update or create several Solr documents 164 | type SolrBatchUpdateQuery struct { 165 | Documents []map[string]interface{} 166 | CommitWithin int 167 | handler string 168 | resultCh chan *SolrResponse 169 | } 170 | 171 | func NewSolrBatchUpdateQuery(documents []map[string]interface{}) *SolrBatchUpdateQuery { 172 | return &SolrBatchUpdateQuery{ 173 | Documents: documents, 174 | handler: "update", 175 | CommitWithin: 0, 176 | resultCh: make(chan *SolrResponse, 1), 177 | } 178 | } 179 | 180 | func NewSolrBatchUpdateQueryCommitWithin(timeInMillis int, documents []map[string]interface{}) *SolrBatchUpdateQuery { 181 | q := NewSolrBatchUpdateQuery(documents) 182 | q.CommitWithin = timeInMillis 183 | return q 184 | } 185 | 186 | func (q *SolrBatchUpdateQuery) Handler() string { 187 | return q.handler 188 | } 189 | 190 | func (q *SolrBatchUpdateQuery) ResultCh() chan *SolrResponse { 191 | return q.resultCh 192 | } 193 | 194 | func (q *SolrBatchUpdateQuery) Wait() *SolrResponse { 195 | return <-q.ResultCh() 196 | } 197 | 198 | func (q *SolrBatchUpdateQuery) GetRows() int { 199 | return 0 200 | } 201 | 202 | func (q *SolrBatchUpdateQuery) GetStart() int { 203 | return 0 204 | } 205 | 206 | func (q *SolrBatchUpdateQuery) Bytes() []byte { 207 | docs := make([]string, len(q.Documents)) 208 | for i, d := range q.Documents { 209 | b, _ := json.Marshal(d) 210 | if q.CommitWithin > 0 { 211 | docs[i] = fmt.Sprintf("\"add\":{\"doc\":%s,\"commitWithin\":%d}", b, q.CommitWithin) 212 | } else { 213 | docs[i] = fmt.Sprintf("\"add\":{\"doc\":%s}", b) 214 | } 215 | } 216 | 217 | buf := strings.Join(docs, ",") 218 | var buffer *bytes.Buffer 219 | if q.CommitWithin > 0 { 220 | buffer = bytes.NewBufferString(fmt.Sprintf("{%s}", buf)) 221 | } else { 222 | buffer = bytes.NewBufferString(fmt.Sprintf("{%s, \"commit\": {}}", buf)) 223 | } 224 | 225 | return buffer.Bytes() 226 | } 227 | 228 | // SolrDeleteQuery represents a query that will remove documents 229 | type SolrDeleteQuery struct { 230 | handler string 231 | match string 232 | resultCh chan *SolrResponse 233 | } 234 | 235 | func NewSolrDeleteQuery(match string) *SolrDeleteQuery { 236 | return &SolrDeleteQuery{ 237 | handler: "update", 238 | match: match, 239 | resultCh: make(chan *SolrResponse, 1), 240 | } 241 | } 242 | 243 | func (q *SolrDeleteQuery) Handler() string { 244 | return q.handler 245 | } 246 | 247 | func (q *SolrDeleteQuery) ResultCh() chan *SolrResponse { 248 | return q.resultCh 249 | } 250 | 251 | func (q *SolrDeleteQuery) Wait() *SolrResponse { 252 | return <-q.ResultCh() 253 | } 254 | 255 | func (q *SolrDeleteQuery) GetRows() int { 256 | return 0 257 | } 258 | 259 | func (q *SolrDeleteQuery) GetStart() int { 260 | return 0 261 | } 262 | 263 | func (q *SolrDeleteQuery) Bytes() []byte { 264 | query := fmt.Sprintf("{\"delete\":{\"query\":%s}, \"commit\": {}}", strconv.Quote(q.match)) 265 | buffer := bytes.NewBufferString(query) 266 | 267 | return buffer.Bytes() 268 | } 269 | 270 | // SolrBatchDeleteQuery represents a query that will remove documents 271 | type SolrBatchDeleteQuery struct { 272 | Ids []string 273 | handler string 274 | resultCh chan *SolrResponse 275 | } 276 | 277 | func NewSolrBatchDeleteQuery(ids []string) *SolrBatchDeleteQuery { 278 | return &SolrBatchDeleteQuery{ 279 | handler: "update", 280 | Ids: ids, 281 | resultCh: make(chan *SolrResponse, 1), 282 | } 283 | } 284 | 285 | func (q *SolrBatchDeleteQuery) Handler() string { 286 | return q.handler 287 | } 288 | 289 | func (q *SolrBatchDeleteQuery) ResultCh() chan *SolrResponse { 290 | return q.resultCh 291 | } 292 | 293 | func (q *SolrBatchDeleteQuery) Wait() *SolrResponse { 294 | return <-q.ResultCh() 295 | } 296 | 297 | func (q *SolrBatchDeleteQuery) GetRows() int { 298 | return 0 299 | } 300 | 301 | func (q *SolrBatchDeleteQuery) GetStart() int { 302 | return 0 303 | } 304 | 305 | func (q *SolrBatchDeleteQuery) Bytes() []byte { 306 | b, _ := json.Marshal(q.Ids) 307 | query := fmt.Sprintf(`"delete":%s`, b) 308 | 309 | buffer := bytes.NewBufferString(fmt.Sprintf(`{%s, "commit": {}}`, query)) 310 | 311 | return buffer.Bytes() 312 | } 313 | -------------------------------------------------------------------------------- /solrquery_test.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | import ( 4 | "bytes" 5 | "encoding/json" 6 | "fmt" 7 | "reflect" 8 | "testing" 9 | ) 10 | 11 | func TestSolrQuery(t *testing.T) { 12 | expectedParams := make(map[string]interface{}) 13 | expectedParams["wt"] = "json" 14 | expectedParams["start"] = float64(0) 15 | expectedParams["rows"] = float64(0) 16 | 17 | expected := make(map[string]interface{}) 18 | expected["query"] = "greeting:你好 AND date:1January2016" 19 | expected["params"] = expectedParams 20 | 21 | var result map[string]interface{} 22 | 23 | query := fmt.Sprintf("greeting:你好 AND date:1January2016") 24 | solrQuery := NewSolrQuery(query, 0, 0, nil, nil, nil, "select") 25 | 26 | err := json.Unmarshal(solrQuery.Bytes(), &result) 27 | if err != nil { 28 | t.Fatal("Unexpected error ", err) 29 | } 30 | 31 | if !reflect.DeepEqual(result, expected) { 32 | t.Error("Result was unexpected ", result) 33 | } 34 | } 35 | 36 | func TestSpatialQuery(t *testing.T) { 37 | expectedParams := make(map[string]interface{}) 38 | expectedParams["wt"] = "json" 39 | expectedParams["fq"] = "{!bbox sfield=latlon}" 40 | expectedParams["d"] = "1.000000" 41 | expectedParams["pt"] = "1.230000,-4.560000" 42 | expectedParams["start"] = float64(0) 43 | expectedParams["rows"] = float64(0) 44 | 45 | expected := make(map[string]interface{}) 46 | expected["query"] = "greeting:你好 AND date:1January2016" 47 | expected["params"] = expectedParams 48 | 49 | var result map[string]interface{} 50 | 51 | query := fmt.Sprintf("greeting:你好 AND date:1January2016") 52 | solrQuery := NewSolrSpatialQuery(query, "bbox", "latlon", 1.23, -4.56, 1, 0, 0, nil, nil, nil, "select") 53 | 54 | err := json.Unmarshal(solrQuery.Bytes(), &result) 55 | if err != nil { 56 | t.Fatal("Unexpected error ", err) 57 | } 58 | 59 | if !reflect.DeepEqual(result, expected) { 60 | t.Error("Result was unexpected ", result) 61 | } 62 | } 63 | 64 | func TestSolrUpdateQuery(t *testing.T) { 65 | expected := []byte(`{"add":{"doc":{"deeper":{"one":"one","two":"two"},"id":"test-id","int_list":[1,2,0,3],"nil":null,"string_list":["你好","Jedná se o delší položka",""]}}, "commit": {}}`) 66 | 67 | deeper := make(map[string]string) 68 | deeper["one"] = "one" 69 | deeper["two"] = "two" 70 | 71 | query := make(map[string]interface{}) 72 | query["id"] = "test-id" 73 | query["string_list"] = []string{"你好", "Jedná se o delší položka", ""} 74 | query["int_list"] = []int{1, 2, 0, 3} 75 | query["deeper"] = deeper 76 | query["nil"] = nil 77 | 78 | solrQuery := NewSolrUpdateQuery(query) 79 | 80 | if bytes.Compare(expected, solrQuery.Bytes()) != 0 { 81 | t.Errorf("Found unexpected query data: %s", solrQuery.Bytes()) 82 | } 83 | } 84 | 85 | func TestSolrBatchUpdateQuery(t *testing.T) { 86 | expected := []byte(`{"add":{"doc":{"deeper":{"one":"one","two":"two"},"id":"test-id","int_list":[1,2,0,3],"nil":null,"string_list":["你好","Jedná se o delší položka",""]}},"add":{"doc":{"deeper":{"one":"one","two":"two"},"id":"test-id2","int_list":[2,4,6,8],"nil":null,"string_list":["你好","Jedná se o delší položka",""]}}, "commit": {}}`) 87 | 88 | deeper := make(map[string]string) 89 | deeper["one"] = "one" 90 | deeper["two"] = "two" 91 | 92 | query := make(map[string]interface{}) 93 | query["id"] = "test-id" 94 | query["string_list"] = []string{"你好", "Jedná se o delší položka", ""} 95 | query["int_list"] = []int{1, 2, 0, 3} 96 | query["deeper"] = deeper 97 | query["nil"] = nil 98 | 99 | nextQuery := make(map[string]interface{}) 100 | nextQuery["id"] = "test-id2" 101 | nextQuery["string_list"] = []string{"你好", "Jedná se o delší položka", ""} 102 | nextQuery["int_list"] = []int{2, 4, 6, 8} 103 | nextQuery["deeper"] = deeper 104 | nextQuery["nil"] = nil 105 | 106 | queries := []map[string]interface{}{query, nextQuery} 107 | solrQuery := NewSolrBatchUpdateQuery(queries) 108 | 109 | if bytes.Compare(expected, solrQuery.Bytes()) != 0 { 110 | t.Errorf("Found unexpected query data: %s", solrQuery.Bytes()) 111 | } 112 | } 113 | 114 | func TestSolrDeleteQuery(t *testing.T) { 115 | expected := []byte("{\"delete\":{\"query\":\"*:*\"}, \"commit\": {}}") 116 | query := NewSolrDeleteQuery("*:*") 117 | 118 | if bytes.Compare(expected, query.Bytes()) != 0 { 119 | t.Errorf("Found unexpected query data: %s", query.Bytes()) 120 | } 121 | 122 | } 123 | 124 | func TestSolrBatchDeleteQuery(t *testing.T) { 125 | expected := []byte(`{"delete":["one","two","three"], "commit": {}}`) 126 | query := NewSolrBatchDeleteQuery([]string{"one", "two", "three"}) 127 | 128 | if bytes.Compare(expected, query.Bytes()) != 0 { 129 | t.Errorf("Found unexpected query data: %s", query.Bytes()) 130 | } 131 | } 132 | -------------------------------------------------------------------------------- /solrresponse.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | import ( 4 | "encoding/json" 5 | "errors" 6 | ) 7 | 8 | var ( 9 | ErrNoResponseHeader = errors.New("Missing response header") 10 | ErrNoDocs = errors.New("No documents found") 11 | ErrBadResponseType = errors.New("Response is an unexpected type") 12 | ErrBadDocs = errors.New("Documents is an unexpected type") 13 | ErrInvalidHeader = errors.New("ResponseHeader was invalid") 14 | ) 15 | 16 | // DocumentCollection represents a collection of solr documents 17 | // and various other metrics 18 | type DocumentCollection struct { 19 | Docs []map[string]interface{} 20 | NumFound int 21 | Start int 22 | } 23 | 24 | // SolrResponse represents a Solr response 25 | type SolrResponse struct { 26 | Facets map[string]interface{} 27 | Response *DocumentCollection 28 | Status int 29 | QTime int 30 | Error error 31 | } 32 | 33 | // PopulateResponse will enumerate the fields of the passed map and create 34 | // a SolrResponse. Only the "responseHeader" field is required. If there 35 | // is a "response" field, it must contain "docs", even if empty. 36 | func PopulateResponse(j map[string]interface{}) (*SolrResponse, error) { 37 | // look for a response element, bail if not present 38 | response_root := j 39 | response := response_root["response"] 40 | 41 | // begin Response creation 42 | r := SolrResponse{} 43 | 44 | // do status & qtime, if possible 45 | r_header, ok := response_root["responseHeader"].(map[string]interface{}) 46 | if !ok { 47 | return nil, ErrNoResponseHeader 48 | } 49 | 50 | if status, ok := r_header["status"]; ok { 51 | r.Status = int(status.(float64)) 52 | } else { 53 | return nil, ErrInvalidHeader 54 | } 55 | 56 | if qtime, ok := r_header["QTime"]; ok { 57 | r.QTime = int(qtime.(float64)) 58 | } else { 59 | return nil, ErrInvalidHeader 60 | } 61 | 62 | // now do docs, if they exist in the response 63 | if response != nil { 64 | responseMap, ok := response.(map[string]interface{}) 65 | if !ok { 66 | return nil, ErrBadResponseType 67 | } 68 | 69 | docs, ok := responseMap["docs"] 70 | if !ok { 71 | return nil, ErrNoDocs 72 | } 73 | 74 | docsSlice, ok := docs.([]interface{}) 75 | if !ok { 76 | return nil, ErrBadDocs 77 | } 78 | 79 | // the total amount of results, irrespective of the amount returned in the response 80 | num_found := int(responseMap["numFound"].(float64)) 81 | 82 | // and the amount actually returned 83 | num_results := len(docsSlice) 84 | 85 | coll := DocumentCollection{} 86 | coll.NumFound = num_found 87 | 88 | ds := make([]map[string]interface{}, 0, num_results) 89 | 90 | for i := 0; i < num_results; i++ { 91 | document, ok := docsSlice[i].(map[string]interface{}) 92 | if ok { 93 | ds = append(ds, document) 94 | } 95 | } 96 | 97 | coll.Docs = ds 98 | r.Response = &coll 99 | } 100 | 101 | // If facets exist, add them as well 102 | facets := response_root["facets"] 103 | if facets != nil { 104 | r.Facets = facets.(map[string]interface{}) 105 | } 106 | 107 | if r.Status >= 400 { 108 | solrError, ok := response_root["error"] 109 | if ok { 110 | errMap, ok := solrError.(map[string]interface{}) 111 | if ok { 112 | solrMsg, ok := errMap["msg"] 113 | if ok { 114 | r.Error = errors.New(solrMsg.(string)) 115 | } 116 | } 117 | } 118 | } 119 | 120 | return &r, nil 121 | } 122 | 123 | // SolrResponseFromHTTPResponse decodes an HTTP (Solr) response 124 | func SolrResponseFromHTTPResponse(b []byte) (*SolrResponse, error) { 125 | var container map[string]interface{} 126 | 127 | err := json.Unmarshal(b, &container) 128 | if err != nil { 129 | return nil, err 130 | } 131 | 132 | resp, err := PopulateResponse(container) 133 | if err != nil { 134 | return nil, err 135 | } 136 | 137 | return resp, nil 138 | } 139 | -------------------------------------------------------------------------------- /solrresponse_test.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestPopulateResponse(t *testing.T) { 8 | data := make(map[string]interface{}) 9 | data["bad"] = true 10 | 11 | _, err := PopulateResponse(data) 12 | if err != ErrNoResponseHeader { 13 | t.Errorf("Expected NoResponseHeaderErr, got %v", err) 14 | } 15 | 16 | data = make(map[string]interface{}) 17 | data["responseHeader"] = make(map[string]interface{}) 18 | 19 | _, err = PopulateResponse(data) 20 | if err != ErrInvalidHeader { 21 | t.Errorf("Expected InvalidHeaderErr, got %v", err) 22 | } 23 | 24 | header := make(map[string]interface{}) 25 | header["status"] = float64(7) 26 | header["QTime"] = float64(4) 27 | 28 | data = make(map[string]interface{}) 29 | data["responseHeader"] = header 30 | 31 | response, err := PopulateResponse(data) 32 | if err != nil { 33 | t.Errorf("Unexpected error: %s", err) 34 | } 35 | 36 | if response.Status != 7 { 37 | t.Errorf("Expected status of 7, got %d", response.Status) 38 | } 39 | 40 | if response.QTime != 4 { 41 | t.Errorf("Expected QTime of 4, got %d", response.QTime) 42 | } 43 | } 44 | 45 | func TestFromHttpResponse(t *testing.T) { 46 | raw := []byte(`{not valid json}`) 47 | _, err := SolrResponseFromHTTPResponse(raw) 48 | if err == nil { 49 | t.Error("Expected json unmarshal error") 50 | } 51 | 52 | raw = []byte(`{"responseHeader": {}}`) 53 | _, err = SolrResponseFromHTTPResponse(raw) 54 | if err != ErrInvalidHeader { 55 | t.Errorf("Expected InvalidHeaderErr, got %s", err) 56 | } 57 | 58 | raw = []byte(`{ 59 | "responseHeader": { 60 | "status": 0, 61 | "QTime": 72, 62 | "params": { 63 | "q": "*:*", 64 | "indent": "true", 65 | "wt": "json", 66 | "_": "1456851532479" 67 | } 68 | }, 69 | "response": { 70 | "numFound": 21, 71 | "start": 0, 72 | "docs": [ 73 | { 74 | "id": "My Id", 75 | "greeting": "你好", 76 | "list": [ 77 | "Jedná se o delší položka" 78 | ], 79 | "_version_": 1525161299814645800 80 | } 81 | ] 82 | } 83 | }`) 84 | 85 | response, err := SolrResponseFromHTTPResponse(raw) 86 | if err != nil { 87 | t.Errorf("Unexpected error %s", err) 88 | } 89 | 90 | if len(response.Response.Docs) != 1 { 91 | t.Errorf("Expected 1 doument, found %d", len(response.Response.Docs)) 92 | } 93 | 94 | raw = []byte(`{ 95 | "responseHeader": { 96 | "status": 0, 97 | "QTime": 72, 98 | "params": { 99 | "q": "*:*", 100 | "indent": "true", 101 | "wt": "json", 102 | "_": "1456851532479" 103 | } 104 | }, 105 | "response": { 106 | "numFound": 0, 107 | "start": 0, 108 | "docs": [] 109 | } 110 | }`) 111 | 112 | response, err = SolrResponseFromHTTPResponse(raw) 113 | if err != nil { 114 | t.Errorf("Unexpected error %s", err) 115 | } 116 | 117 | if len(response.Response.Docs) != 0 { 118 | t.Errorf("Expected 0 doument, found %d", len(response.Response.Docs)) 119 | } 120 | 121 | } 122 | -------------------------------------------------------------------------------- /workerpool.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | import ( 4 | "errors" 5 | "sync" 6 | "time" 7 | 8 | "github.com/wirelessregistry/glog" 9 | ) 10 | 11 | var ( 12 | ErrPoolNotRunning = errors.New("Pool is not running.") 13 | ErrPoolRunning = errors.New("Pool is running.") 14 | ErrNoActiveWorkers = errors.New("Pool has no active workers.") 15 | ErrTimeout = errors.New("Host Timeout") 16 | ) 17 | 18 | type worker struct { 19 | parent *Pool 20 | client SolrClient 21 | jobCh <-chan SolrJob 22 | dieCh <-chan struct{} 23 | sigDeathCh chan<- struct{} 24 | timeout int 25 | } 26 | 27 | func newWorker(parent *Pool, jCh <-chan SolrJob, dCh <-chan struct{}, sDeathCh chan<- struct{}, client SolrClient, timeout int) *worker { 28 | return &worker{ 29 | parent: parent, 30 | jobCh: jCh, 31 | dieCh: dCh, 32 | sigDeathCh: sDeathCh, 33 | client: client, 34 | timeout: timeout, 35 | } 36 | } 37 | 38 | func (w *worker) work() { 39 | var resp *SolrResponse 40 | var hostReachable <-chan time.Time 41 | 42 | timeout := false 43 | jCh := w.jobCh 44 | hostReachable = nil 45 | 46 | for { 47 | if timeout { 48 | jCh = nil 49 | hostReachable = time.After(time.Second * time.Duration(w.timeout)) 50 | } else { 51 | jCh = w.jobCh 52 | hostReachable = nil 53 | } 54 | 55 | select { 56 | case <-w.dieCh: 57 | w.sigDeathCh <- struct{}{} 58 | return 59 | 60 | case job := <-jCh: 61 | resp, timeout = w.client.Execute(job) 62 | if timeout { 63 | glog.Warning("SolrWorker.timeout received.") 64 | resp.Error = ErrTimeout 65 | } 66 | 67 | job.ResultCh() <- resp 68 | case <-hostReachable: 69 | if glog.V(2) { 70 | glog.Info("Trying to reconnect to host...") 71 | } 72 | timeout = w.hostOffline() 73 | glog.Warningf("SolrWorker.hostReachable = %v.", timeout) 74 | } 75 | } 76 | } 77 | 78 | func (w *worker) hostOffline() bool { 79 | return w.client.TestConnection() 80 | } 81 | 82 | // Pool holds all the data about our worker pool 83 | type Pool struct { 84 | // nWorkers specifies the total number of workers this pool should have 85 | nWorkersPerClient int 86 | 87 | // bufferLen specifies the number of jobs that can be in queue without blocking 88 | bufferLen int 89 | 90 | clients []SolrClient 91 | 92 | timeout int 93 | jobCh chan SolrJob 94 | dieCh chan struct{} 95 | lock sync.Mutex 96 | } 97 | 98 | // NewPool will create a Pool structure with an array of Solr servers. 99 | // It will create numWorkers per Solr server, and allow bufLen jobs 100 | // before the workers start blocking. 101 | func NewPool(clients []SolrClient, numWorkersPerClient, bufLen, timeout int) *Pool { 102 | p := &Pool{} 103 | p.clients = clients 104 | p.nWorkersPerClient = numWorkersPerClient 105 | p.bufferLen = bufLen 106 | p.timeout = timeout 107 | return p 108 | } 109 | 110 | // Run will create a goroutine for each worker, and a master goroutine 111 | // to control those workers. A channel to close down the pool will be 112 | // returned to the caller. 113 | func (p *Pool) Run() (<-chan struct{}, error) { 114 | p.lock.Lock() 115 | defer p.lock.Unlock() 116 | 117 | if p.jobCh != nil { 118 | return nil, ErrPoolRunning 119 | } 120 | 121 | glog.Infof("SolrPool.Run() with %v worker(s).", p.nWorkersPerClient) 122 | 123 | p.jobCh = make(chan SolrJob, p.bufferLen) 124 | p.dieCh = make(chan struct{}, 1) 125 | 126 | sigPoolDeathCh := make(chan struct{}, 1) 127 | collectWorkersDeathCh := make(chan struct{}, p.nWorkersPerClient) 128 | dieChs := make([]chan struct{}, 0, p.nWorkersPerClient) 129 | 130 | for _, client := range p.clients { 131 | for i := 0; i < p.nWorkersPerClient; i++ { 132 | dieCh := make(chan struct{}, 1) 133 | dieChs = append(dieChs, dieCh) 134 | 135 | w := newWorker(p, p.jobCh, dieCh, collectWorkersDeathCh, client, p.timeout) 136 | go w.work() 137 | } 138 | } 139 | 140 | go master(p.dieCh, dieChs, sigPoolDeathCh, collectWorkersDeathCh) 141 | 142 | return sigPoolDeathCh, nil 143 | } 144 | 145 | // Submit will enter a job into the queue for the worker pool 146 | func (p *Pool) Submit(s SolrJob) error { 147 | p.lock.Lock() 148 | defer p.lock.Unlock() 149 | 150 | if p.jobCh == nil { 151 | return ErrPoolNotRunning 152 | } 153 | 154 | p.jobCh <- s 155 | return nil 156 | } 157 | 158 | // Stop will gracefully stop processing jobs and shutdown the workers 159 | func (p *Pool) Stop() { 160 | p.lock.Lock() 161 | defer p.lock.Unlock() 162 | 163 | if p.jobCh == nil { 164 | return 165 | } 166 | 167 | p.dieCh <- struct{}{} 168 | p.jobCh = nil 169 | } 170 | 171 | func master(dieCh chan struct{}, dieChs []chan struct{}, sigPoolDeathCh chan<- struct{}, collectWorkersDeathCh <-chan struct{}) { 172 | workersLeft := len(dieChs) 173 | 174 | for { 175 | if workersLeft == 0 { 176 | sigPoolDeathCh <- struct{}{} 177 | return 178 | } 179 | 180 | select { 181 | case <-dieCh: 182 | for i, _ := range dieChs { 183 | dieChs[i] <- struct{}{} 184 | } 185 | 186 | case <-collectWorkersDeathCh: 187 | workersLeft-- 188 | } 189 | } 190 | } 191 | -------------------------------------------------------------------------------- /workerpool_test.go: -------------------------------------------------------------------------------- 1 | package gora 2 | 3 | import ( 4 | "strconv" 5 | "sync" 6 | "testing" 7 | "time" 8 | ) 9 | 10 | func TestPoolNoSubmission(t *testing.T) { 11 | ch := make(chan struct{}, 1) 12 | client := MockClientConstructor("0.0.0.0", "", ch) 13 | p := NewPool([]SolrClient{client}, 10, 1, 1) 14 | err := p.Submit(nil) 15 | if err != ErrPoolNotRunning { 16 | t.Fatalf("Expected errPoolNotRunning") 17 | } 18 | 19 | sigCh, _ := p.Run() 20 | p.Stop() 21 | 22 | sigCh2, _ := p.Run() 23 | _, err = p.Run() 24 | if err != ErrPoolRunning { 25 | t.Errorf("Expected %v. Got %v.", ErrPoolRunning, err) 26 | } 27 | 28 | p.Stop() 29 | 30 | <-sigCh 31 | <-sigCh2 32 | } 33 | 34 | func TestPoolMock(t *testing.T) { 35 | testPoolMock(t, 1, 0) 36 | testPoolMock(t, 10, 0) 37 | testPoolMock(t, 50, 1) 38 | } 39 | 40 | func testPoolMock(t *testing.T, nw int, bl int) { 41 | var wg sync.WaitGroup 42 | 43 | ch := make(chan struct{}, 1) 44 | client := MockClientConstructor("0.0.0.0", "", ch) 45 | p := NewPool([]SolrClient{client}, nw, bl, 1) 46 | sig, err := p.Run() 47 | if err != nil { 48 | t.Fatal("Unexpected error ", err) 49 | } 50 | 51 | for i := 0; i < 9999; i++ { 52 | job := NewMockSolrJob([]byte(strconv.Itoa(i))) 53 | 54 | go func(j SolrJob) { 55 | p.Submit(j) 56 | }(job) 57 | 58 | wg.Add(1) 59 | go func(j SolrJob, expected int) { 60 | j.Wait() 61 | wg.Done() 62 | }(job, i) 63 | } 64 | 65 | wg.Wait() 66 | p.Stop() 67 | <-sig 68 | } 69 | 70 | func TestPoolFaultyWorker(t *testing.T) { 71 | ch := make(chan struct{}, 1) 72 | client := &MockSolrClient{ 73 | Timeout: true, 74 | CloseCh: ch, 75 | Response: &SolrResponse{}, 76 | } 77 | 78 | p := NewPool([]SolrClient{client}, 1, 0, 1) 79 | sig, _ := p.Run() 80 | 81 | job := NewMockSolrJob([]byte("dead")) 82 | p.Submit(job) 83 | 84 | resp := job.Wait() 85 | if resp.Error != ErrTimeout { 86 | t.Errorf("Expected %v. Got %v.", ErrTimeout, resp.Error) 87 | } 88 | 89 | // Must allow wait time for 3 connection fails 90 | waitTime := time.After(time.Second * 5) 91 | select { 92 | case <-ch: 93 | case <-waitTime: 94 | t.Error("Got timeout waiting for mock client") 95 | } 96 | p.Stop() 97 | <-sig 98 | } 99 | --------------------------------------------------------------------------------